VmapGeneratedPlumbing.h 1.7 MB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771177217731774177517761777177817791780178117821783178417851786178717881789179017911792179317941795179617971798179918001801180218031804180518061807180818091810181118121813181418151816181718181819182018211822182318241825182618271828182918301831183218331834183518361837183818391840184118421843184418451846184718481849185018511852185318541855185618571858185918601861186218631864186518661867186818691870187118721873187418751876187718781879188018811882188318841885188618871888188918901891189218931894189518961897189818991900190119021903190419051906190719081909191019111912191319141915191619171918191919201921192219231924192519261927192819291930193119321933193419351936193719381939194019411942194319441945194619471948194919501951195219531954195519561957195819591960196119621963196419651966196719681969197019711972197319741975197619771978197919801981198219831984198519861987198819891990199119921993199419951996199719981999200020012002200320042005200620072008200920102011201220132014201520162017201820192020202120222023202420252026202720282029203020312032203320342035203620372038203920402041204220432044204520462047204820492050205120522053205420552056205720582059206020612062206320642065206620672068206920702071207220732074207520762077207820792080208120822083208420852086208720882089209020912092209320942095209620972098209921002101210221032104210521062107210821092110211121122113211421152116211721182119212021212122212321242125212621272128212921302131213221332134213521362137213821392140214121422143214421452146214721482149215021512152215321542155215621572158215921602161216221632164216521662167216821692170217121722173217421752176217721782179218021812182218321842185218621872188218921902191219221932194219521962197219821992200220122022203220422052206220722082209221022112212221322142215221622172218221922202221222222232224222522262227222822292230223122322233223422352236223722382239224022412242224322442245224622472248224922502251225222532254225522562257225822592260226122622263226422652266226722682269227022712272227322742275227622772278227922802281228222832284228522862287228822892290229122922293229422952296229722982299230023012302230323042305230623072308230923102311231223132314231523162317231823192320232123222323232423252326232723282329233023312332233323342335233623372338233923402341234223432344234523462347234823492350235123522353235423552356235723582359236023612362236323642365236623672368236923702371237223732374237523762377237823792380238123822383238423852386238723882389239023912392239323942395239623972398239924002401240224032404240524062407240824092410241124122413241424152416241724182419242024212422242324242425242624272428242924302431243224332434243524362437243824392440244124422443244424452446244724482449245024512452245324542455245624572458245924602461246224632464246524662467246824692470247124722473247424752476247724782479248024812482248324842485248624872488248924902491249224932494249524962497249824992500250125022503250425052506250725082509251025112512251325142515251625172518251925202521252225232524252525262527252825292530253125322533253425352536253725382539254025412542254325442545254625472548254925502551255225532554255525562557255825592560256125622563256425652566256725682569257025712572257325742575257625772578257925802581258225832584258525862587258825892590259125922593259425952596259725982599260026012602260326042605260626072608260926102611261226132614261526162617261826192620262126222623262426252626262726282629263026312632263326342635263626372638263926402641264226432644264526462647264826492650265126522653265426552656265726582659266026612662266326642665266626672668266926702671267226732674267526762677267826792680268126822683268426852686268726882689269026912692269326942695269626972698269927002701270227032704270527062707270827092710271127122713271427152716271727182719272027212722272327242725272627272728272927302731273227332734273527362737273827392740274127422743274427452746274727482749275027512752275327542755275627572758275927602761276227632764276527662767276827692770277127722773277427752776277727782779278027812782278327842785278627872788278927902791279227932794279527962797279827992800280128022803280428052806280728082809281028112812281328142815281628172818281928202821282228232824282528262827282828292830283128322833283428352836283728382839284028412842284328442845284628472848284928502851285228532854285528562857285828592860286128622863286428652866286728682869287028712872287328742875287628772878287928802881288228832884288528862887288828892890289128922893289428952896289728982899290029012902290329042905290629072908290929102911291229132914291529162917291829192920292129222923292429252926292729282929293029312932293329342935293629372938293929402941294229432944294529462947294829492950295129522953295429552956295729582959296029612962296329642965296629672968296929702971297229732974297529762977297829792980298129822983298429852986298729882989299029912992299329942995299629972998299930003001300230033004300530063007300830093010301130123013301430153016301730183019302030213022302330243025302630273028302930303031303230333034303530363037303830393040304130423043304430453046304730483049305030513052305330543055305630573058305930603061306230633064306530663067306830693070307130723073307430753076307730783079308030813082308330843085308630873088308930903091309230933094309530963097309830993100310131023103310431053106310731083109311031113112311331143115311631173118311931203121312231233124312531263127312831293130313131323133313431353136313731383139314031413142314331443145314631473148314931503151315231533154315531563157315831593160316131623163316431653166316731683169317031713172317331743175317631773178317931803181318231833184318531863187318831893190319131923193319431953196319731983199320032013202320332043205320632073208320932103211321232133214321532163217321832193220322132223223322432253226322732283229323032313232323332343235323632373238323932403241324232433244324532463247324832493250325132523253325432553256325732583259326032613262326332643265326632673268326932703271327232733274327532763277327832793280328132823283328432853286328732883289329032913292329332943295329632973298329933003301330233033304330533063307330833093310331133123313331433153316331733183319332033213322332333243325332633273328332933303331333233333334333533363337333833393340334133423343334433453346334733483349335033513352335333543355335633573358335933603361336233633364336533663367336833693370337133723373337433753376337733783379338033813382338333843385338633873388338933903391339233933394339533963397339833993400340134023403340434053406340734083409341034113412341334143415341634173418341934203421342234233424342534263427342834293430343134323433343434353436343734383439344034413442344334443445344634473448344934503451345234533454345534563457345834593460346134623463346434653466346734683469347034713472347334743475347634773478347934803481348234833484348534863487348834893490349134923493349434953496349734983499350035013502350335043505350635073508350935103511351235133514351535163517351835193520352135223523352435253526352735283529353035313532353335343535353635373538353935403541354235433544354535463547354835493550355135523553355435553556355735583559356035613562356335643565356635673568356935703571357235733574357535763577357835793580358135823583358435853586358735883589359035913592359335943595359635973598359936003601360236033604360536063607360836093610361136123613361436153616361736183619362036213622362336243625362636273628362936303631363236333634363536363637363836393640364136423643364436453646364736483649365036513652365336543655365636573658365936603661366236633664366536663667366836693670367136723673367436753676367736783679368036813682368336843685368636873688368936903691369236933694369536963697369836993700370137023703370437053706370737083709371037113712371337143715371637173718371937203721372237233724372537263727372837293730373137323733373437353736373737383739374037413742374337443745374637473748374937503751375237533754375537563757375837593760376137623763376437653766376737683769377037713772377337743775377637773778377937803781378237833784378537863787378837893790379137923793379437953796379737983799380038013802380338043805380638073808380938103811381238133814381538163817381838193820382138223823382438253826382738283829383038313832383338343835383638373838383938403841384238433844384538463847384838493850385138523853385438553856385738583859386038613862386338643865386638673868386938703871387238733874387538763877387838793880388138823883388438853886388738883889389038913892389338943895389638973898389939003901390239033904390539063907390839093910391139123913391439153916391739183919392039213922392339243925392639273928392939303931393239333934393539363937393839393940394139423943394439453946394739483949395039513952395339543955395639573958395939603961396239633964396539663967396839693970397139723973397439753976397739783979398039813982398339843985398639873988398939903991399239933994399539963997399839994000400140024003400440054006400740084009401040114012401340144015401640174018401940204021402240234024402540264027402840294030403140324033403440354036403740384039404040414042404340444045404640474048404940504051405240534054405540564057405840594060406140624063406440654066406740684069407040714072407340744075407640774078407940804081408240834084408540864087408840894090409140924093409440954096409740984099410041014102410341044105410641074108410941104111411241134114411541164117411841194120412141224123412441254126412741284129413041314132413341344135413641374138413941404141414241434144414541464147414841494150415141524153415441554156415741584159416041614162416341644165416641674168416941704171417241734174417541764177417841794180418141824183418441854186418741884189419041914192419341944195419641974198419942004201420242034204420542064207420842094210421142124213421442154216421742184219422042214222422342244225422642274228422942304231423242334234423542364237423842394240424142424243424442454246424742484249425042514252425342544255425642574258425942604261426242634264426542664267426842694270427142724273427442754276427742784279428042814282428342844285428642874288428942904291429242934294429542964297429842994300430143024303430443054306430743084309431043114312431343144315431643174318431943204321432243234324432543264327432843294330433143324333433443354336433743384339434043414342434343444345434643474348434943504351435243534354435543564357435843594360436143624363436443654366436743684369437043714372437343744375437643774378437943804381438243834384438543864387438843894390439143924393439443954396439743984399440044014402440344044405440644074408440944104411441244134414441544164417441844194420442144224423442444254426442744284429443044314432443344344435443644374438443944404441444244434444444544464447444844494450445144524453445444554456445744584459446044614462446344644465446644674468446944704471447244734474447544764477447844794480448144824483448444854486448744884489449044914492449344944495449644974498449945004501450245034504450545064507450845094510451145124513451445154516451745184519452045214522452345244525452645274528452945304531453245334534453545364537453845394540454145424543454445454546454745484549455045514552455345544555455645574558455945604561456245634564456545664567456845694570457145724573457445754576457745784579458045814582458345844585458645874588458945904591459245934594459545964597459845994600460146024603460446054606460746084609461046114612461346144615461646174618461946204621462246234624462546264627462846294630463146324633463446354636463746384639464046414642464346444645464646474648464946504651465246534654465546564657465846594660466146624663466446654666466746684669467046714672467346744675467646774678467946804681468246834684468546864687468846894690469146924693469446954696469746984699470047014702470347044705470647074708470947104711471247134714471547164717471847194720472147224723472447254726472747284729473047314732473347344735473647374738473947404741474247434744474547464747474847494750475147524753475447554756475747584759476047614762476347644765476647674768476947704771477247734774477547764777477847794780478147824783478447854786478747884789479047914792479347944795479647974798479948004801480248034804480548064807480848094810481148124813481448154816481748184819482048214822482348244825482648274828482948304831483248334834483548364837483848394840484148424843484448454846484748484849485048514852485348544855485648574858485948604861486248634864486548664867486848694870487148724873487448754876487748784879488048814882488348844885488648874888488948904891489248934894489548964897489848994900490149024903490449054906490749084909491049114912491349144915491649174918491949204921492249234924492549264927492849294930493149324933493449354936493749384939494049414942494349444945494649474948494949504951495249534954495549564957495849594960496149624963496449654966496749684969497049714972497349744975497649774978497949804981498249834984498549864987498849894990499149924993499449954996499749984999500050015002500350045005500650075008500950105011501250135014501550165017501850195020502150225023502450255026502750285029503050315032503350345035503650375038503950405041504250435044504550465047504850495050505150525053505450555056505750585059506050615062506350645065506650675068506950705071507250735074507550765077507850795080508150825083508450855086508750885089509050915092509350945095509650975098509951005101510251035104510551065107510851095110511151125113511451155116511751185119512051215122512351245125512651275128512951305131513251335134513551365137513851395140514151425143514451455146514751485149515051515152515351545155515651575158515951605161516251635164516551665167516851695170517151725173517451755176517751785179518051815182518351845185518651875188518951905191519251935194519551965197519851995200520152025203520452055206520752085209521052115212521352145215521652175218521952205221522252235224522552265227522852295230523152325233523452355236523752385239524052415242524352445245524652475248524952505251525252535254525552565257525852595260526152625263526452655266526752685269527052715272527352745275527652775278527952805281528252835284528552865287528852895290529152925293529452955296529752985299530053015302530353045305530653075308530953105311531253135314531553165317531853195320532153225323532453255326532753285329533053315332533353345335533653375338533953405341534253435344534553465347534853495350535153525353535453555356535753585359536053615362536353645365536653675368536953705371537253735374537553765377537853795380538153825383538453855386538753885389539053915392539353945395539653975398539954005401540254035404540554065407540854095410541154125413541454155416541754185419542054215422542354245425542654275428542954305431543254335434543554365437543854395440544154425443544454455446544754485449545054515452545354545455545654575458545954605461546254635464546554665467546854695470547154725473547454755476547754785479548054815482548354845485548654875488548954905491549254935494549554965497549854995500550155025503550455055506550755085509551055115512551355145515551655175518551955205521552255235524552555265527552855295530553155325533553455355536553755385539554055415542554355445545554655475548554955505551555255535554555555565557555855595560556155625563556455655566556755685569557055715572557355745575557655775578557955805581558255835584558555865587558855895590559155925593559455955596559755985599560056015602560356045605560656075608560956105611561256135614561556165617561856195620562156225623562456255626562756285629563056315632563356345635563656375638563956405641564256435644564556465647564856495650565156525653565456555656565756585659566056615662566356645665566656675668566956705671567256735674567556765677567856795680568156825683568456855686568756885689569056915692569356945695569656975698569957005701570257035704570557065707570857095710571157125713571457155716571757185719572057215722572357245725572657275728572957305731573257335734573557365737573857395740574157425743574457455746574757485749575057515752575357545755575657575758575957605761576257635764576557665767576857695770577157725773577457755776577757785779578057815782578357845785578657875788578957905791579257935794579557965797579857995800580158025803580458055806580758085809581058115812581358145815581658175818581958205821582258235824582558265827582858295830583158325833583458355836583758385839584058415842584358445845584658475848584958505851585258535854585558565857585858595860586158625863586458655866586758685869587058715872587358745875587658775878587958805881588258835884588558865887588858895890589158925893589458955896589758985899590059015902590359045905590659075908590959105911591259135914591559165917591859195920592159225923592459255926592759285929593059315932593359345935593659375938593959405941594259435944594559465947594859495950595159525953595459555956595759585959596059615962596359645965596659675968596959705971597259735974597559765977597859795980598159825983598459855986598759885989599059915992599359945995599659975998599960006001600260036004600560066007600860096010601160126013601460156016601760186019602060216022602360246025602660276028602960306031603260336034603560366037603860396040604160426043604460456046604760486049605060516052605360546055605660576058605960606061606260636064606560666067606860696070607160726073607460756076607760786079608060816082608360846085608660876088608960906091609260936094609560966097609860996100610161026103610461056106610761086109611061116112611361146115611661176118611961206121612261236124612561266127612861296130613161326133613461356136613761386139614061416142614361446145614661476148614961506151615261536154615561566157615861596160616161626163616461656166616761686169617061716172617361746175617661776178617961806181618261836184618561866187618861896190619161926193619461956196619761986199620062016202620362046205620662076208620962106211621262136214621562166217621862196220622162226223622462256226622762286229623062316232623362346235623662376238623962406241624262436244624562466247624862496250625162526253625462556256625762586259626062616262626362646265626662676268626962706271627262736274627562766277627862796280628162826283628462856286628762886289629062916292629362946295629662976298629963006301630263036304630563066307630863096310631163126313631463156316631763186319632063216322632363246325632663276328632963306331633263336334633563366337633863396340634163426343634463456346634763486349635063516352635363546355635663576358635963606361636263636364636563666367636863696370637163726373637463756376637763786379638063816382638363846385638663876388638963906391639263936394639563966397639863996400640164026403640464056406640764086409641064116412641364146415641664176418641964206421642264236424642564266427642864296430643164326433643464356436643764386439644064416442644364446445644664476448644964506451645264536454645564566457645864596460646164626463646464656466646764686469647064716472647364746475647664776478647964806481648264836484648564866487648864896490649164926493649464956496649764986499650065016502650365046505650665076508650965106511651265136514651565166517651865196520652165226523652465256526652765286529653065316532653365346535653665376538653965406541654265436544654565466547654865496550655165526553655465556556655765586559656065616562656365646565656665676568656965706571657265736574657565766577657865796580658165826583658465856586658765886589659065916592659365946595659665976598659966006601660266036604660566066607660866096610661166126613661466156616661766186619662066216622662366246625662666276628662966306631663266336634663566366637663866396640664166426643664466456646664766486649665066516652665366546655665666576658665966606661666266636664666566666667666866696670667166726673667466756676667766786679668066816682668366846685668666876688668966906691669266936694669566966697669866996700670167026703670467056706670767086709671067116712671367146715671667176718671967206721672267236724672567266727672867296730673167326733673467356736673767386739674067416742674367446745674667476748674967506751675267536754675567566757675867596760676167626763676467656766676767686769677067716772677367746775677667776778677967806781678267836784678567866787678867896790679167926793679467956796679767986799680068016802680368046805680668076808680968106811681268136814681568166817681868196820682168226823682468256826682768286829683068316832683368346835683668376838683968406841684268436844684568466847684868496850685168526853685468556856685768586859686068616862686368646865686668676868686968706871687268736874687568766877687868796880688168826883688468856886688768886889689068916892689368946895689668976898689969006901690269036904690569066907690869096910691169126913691469156916691769186919692069216922692369246925692669276928692969306931693269336934693569366937693869396940694169426943694469456946694769486949695069516952695369546955695669576958695969606961696269636964696569666967696869696970697169726973697469756976697769786979698069816982698369846985698669876988698969906991699269936994699569966997699869997000700170027003700470057006700770087009701070117012701370147015701670177018701970207021702270237024702570267027702870297030703170327033703470357036703770387039704070417042704370447045704670477048704970507051705270537054705570567057705870597060706170627063706470657066706770687069707070717072707370747075707670777078707970807081708270837084708570867087708870897090709170927093709470957096709770987099710071017102710371047105710671077108710971107111711271137114711571167117711871197120712171227123712471257126712771287129713071317132713371347135713671377138713971407141714271437144714571467147714871497150715171527153715471557156715771587159716071617162716371647165716671677168716971707171717271737174717571767177717871797180718171827183718471857186718771887189719071917192719371947195719671977198719972007201720272037204720572067207720872097210721172127213721472157216721772187219722072217222722372247225722672277228722972307231723272337234723572367237723872397240724172427243724472457246724772487249725072517252725372547255725672577258725972607261726272637264726572667267726872697270727172727273727472757276727772787279728072817282728372847285728672877288728972907291729272937294729572967297729872997300730173027303730473057306730773087309731073117312731373147315731673177318731973207321732273237324732573267327732873297330733173327333733473357336733773387339734073417342734373447345734673477348734973507351735273537354735573567357735873597360736173627363736473657366736773687369737073717372737373747375737673777378737973807381738273837384738573867387738873897390739173927393739473957396739773987399740074017402740374047405740674077408740974107411741274137414741574167417741874197420742174227423742474257426742774287429743074317432743374347435743674377438743974407441744274437444744574467447744874497450745174527453745474557456745774587459746074617462746374647465746674677468746974707471747274737474747574767477747874797480748174827483748474857486748774887489749074917492749374947495749674977498749975007501750275037504750575067507750875097510751175127513751475157516751775187519752075217522752375247525752675277528752975307531753275337534753575367537753875397540754175427543754475457546754775487549755075517552755375547555755675577558755975607561756275637564756575667567756875697570757175727573757475757576757775787579758075817582758375847585758675877588758975907591759275937594759575967597759875997600760176027603760476057606760776087609761076117612761376147615761676177618761976207621762276237624762576267627762876297630763176327633763476357636763776387639764076417642764376447645764676477648764976507651765276537654765576567657765876597660766176627663766476657666766776687669767076717672767376747675767676777678767976807681768276837684768576867687768876897690769176927693769476957696769776987699770077017702770377047705770677077708770977107711771277137714771577167717771877197720772177227723772477257726772777287729773077317732773377347735773677377738773977407741774277437744774577467747774877497750775177527753775477557756775777587759776077617762776377647765776677677768776977707771777277737774777577767777777877797780778177827783778477857786778777887789779077917792779377947795779677977798779978007801780278037804780578067807780878097810781178127813781478157816781778187819782078217822782378247825782678277828782978307831783278337834783578367837783878397840784178427843784478457846784778487849785078517852785378547855785678577858785978607861786278637864786578667867786878697870787178727873787478757876787778787879788078817882788378847885788678877888788978907891789278937894789578967897789878997900790179027903790479057906790779087909791079117912791379147915791679177918791979207921792279237924792579267927792879297930793179327933793479357936793779387939794079417942794379447945794679477948794979507951795279537954795579567957795879597960796179627963796479657966796779687969797079717972797379747975797679777978797979807981798279837984798579867987798879897990799179927993799479957996799779987999800080018002800380048005800680078008800980108011801280138014801580168017801880198020802180228023802480258026802780288029803080318032803380348035803680378038803980408041804280438044804580468047804880498050805180528053805480558056805780588059806080618062806380648065806680678068806980708071807280738074807580768077807880798080808180828083808480858086808780888089809080918092809380948095809680978098809981008101810281038104810581068107810881098110811181128113811481158116811781188119812081218122812381248125812681278128812981308131813281338134813581368137813881398140814181428143814481458146814781488149815081518152815381548155815681578158815981608161816281638164816581668167816881698170817181728173817481758176817781788179818081818182818381848185818681878188818981908191819281938194819581968197819881998200820182028203820482058206820782088209821082118212821382148215821682178218821982208221822282238224822582268227822882298230823182328233823482358236823782388239824082418242824382448245824682478248824982508251825282538254825582568257825882598260826182628263826482658266826782688269827082718272827382748275827682778278827982808281828282838284828582868287828882898290829182928293829482958296829782988299830083018302830383048305830683078308830983108311831283138314831583168317831883198320832183228323832483258326832783288329833083318332833383348335833683378338833983408341834283438344834583468347834883498350835183528353835483558356835783588359836083618362836383648365836683678368836983708371837283738374837583768377837883798380838183828383838483858386838783888389839083918392839383948395839683978398839984008401840284038404840584068407840884098410841184128413841484158416841784188419842084218422842384248425842684278428842984308431843284338434843584368437843884398440844184428443844484458446844784488449845084518452845384548455845684578458845984608461846284638464846584668467846884698470847184728473847484758476847784788479848084818482848384848485848684878488848984908491849284938494849584968497849884998500850185028503850485058506850785088509851085118512851385148515851685178518851985208521852285238524852585268527852885298530853185328533853485358536853785388539854085418542854385448545854685478548854985508551855285538554855585568557855885598560856185628563856485658566856785688569857085718572857385748575857685778578857985808581858285838584858585868587858885898590859185928593859485958596859785988599860086018602860386048605860686078608860986108611861286138614861586168617861886198620862186228623862486258626862786288629863086318632863386348635863686378638863986408641864286438644864586468647864886498650865186528653865486558656865786588659866086618662866386648665866686678668866986708671867286738674867586768677867886798680868186828683868486858686868786888689869086918692869386948695869686978698869987008701870287038704870587068707870887098710871187128713871487158716871787188719872087218722872387248725872687278728872987308731873287338734873587368737873887398740874187428743874487458746874787488749875087518752875387548755875687578758875987608761876287638764876587668767876887698770877187728773877487758776877787788779878087818782878387848785878687878788878987908791879287938794879587968797879887998800880188028803880488058806880788088809881088118812881388148815881688178818881988208821882288238824882588268827882888298830883188328833883488358836883788388839884088418842884388448845884688478848884988508851885288538854885588568857885888598860886188628863886488658866886788688869887088718872887388748875887688778878887988808881888288838884888588868887888888898890889188928893889488958896889788988899890089018902890389048905890689078908890989108911891289138914891589168917891889198920892189228923892489258926892789288929893089318932893389348935893689378938893989408941894289438944894589468947894889498950895189528953895489558956895789588959896089618962896389648965896689678968896989708971897289738974897589768977897889798980898189828983898489858986898789888989899089918992899389948995899689978998899990009001900290039004900590069007900890099010901190129013901490159016901790189019902090219022902390249025902690279028902990309031903290339034903590369037903890399040904190429043904490459046904790489049905090519052905390549055905690579058905990609061906290639064906590669067906890699070907190729073907490759076907790789079908090819082908390849085908690879088908990909091909290939094909590969097909890999100910191029103910491059106910791089109911091119112911391149115911691179118911991209121912291239124912591269127912891299130913191329133913491359136913791389139914091419142914391449145914691479148914991509151915291539154915591569157915891599160916191629163916491659166916791689169917091719172917391749175917691779178917991809181918291839184918591869187918891899190919191929193919491959196919791989199920092019202920392049205920692079208920992109211921292139214921592169217921892199220922192229223922492259226922792289229923092319232923392349235923692379238923992409241924292439244924592469247924892499250925192529253925492559256925792589259926092619262926392649265926692679268926992709271927292739274927592769277927892799280928192829283928492859286928792889289929092919292929392949295929692979298929993009301930293039304930593069307930893099310931193129313931493159316931793189319932093219322932393249325932693279328932993309331933293339334933593369337933893399340934193429343934493459346934793489349935093519352935393549355935693579358935993609361936293639364936593669367936893699370937193729373937493759376937793789379938093819382938393849385938693879388938993909391939293939394939593969397939893999400940194029403940494059406940794089409941094119412941394149415941694179418941994209421942294239424942594269427942894299430943194329433943494359436943794389439944094419442944394449445944694479448944994509451945294539454945594569457945894599460946194629463946494659466946794689469947094719472947394749475947694779478947994809481948294839484948594869487948894899490949194929493949494959496949794989499950095019502950395049505950695079508950995109511951295139514951595169517951895199520952195229523952495259526952795289529953095319532953395349535953695379538953995409541954295439544954595469547954895499550955195529553955495559556955795589559956095619562956395649565956695679568956995709571957295739574957595769577957895799580958195829583958495859586958795889589959095919592959395949595959695979598959996009601960296039604960596069607960896099610961196129613961496159616961796189619962096219622962396249625962696279628962996309631963296339634963596369637963896399640964196429643964496459646964796489649965096519652965396549655965696579658965996609661966296639664966596669667966896699670967196729673967496759676967796789679968096819682968396849685968696879688968996909691969296939694969596969697969896999700970197029703970497059706970797089709971097119712971397149715971697179718971997209721972297239724972597269727972897299730973197329733973497359736973797389739974097419742974397449745974697479748974997509751975297539754975597569757975897599760976197629763976497659766976797689769977097719772977397749775977697779778977997809781978297839784978597869787978897899790979197929793979497959796979797989799980098019802980398049805980698079808980998109811981298139814981598169817981898199820982198229823982498259826982798289829983098319832983398349835983698379838983998409841984298439844984598469847984898499850985198529853985498559856985798589859986098619862986398649865986698679868986998709871987298739874987598769877987898799880988198829883988498859886988798889889989098919892989398949895989698979898989999009901990299039904990599069907990899099910991199129913991499159916991799189919992099219922992399249925992699279928992999309931993299339934993599369937993899399940994199429943994499459946994799489949995099519952995399549955995699579958995999609961996299639964996599669967996899699970997199729973997499759976997799789979998099819982998399849985998699879988998999909991999299939994999599969997999899991000010001100021000310004100051000610007100081000910010100111001210013100141001510016100171001810019100201002110022100231002410025100261002710028100291003010031100321003310034100351003610037100381003910040100411004210043100441004510046100471004810049100501005110052100531005410055100561005710058100591006010061100621006310064100651006610067100681006910070100711007210073100741007510076100771007810079100801008110082100831008410085100861008710088100891009010091100921009310094100951009610097100981009910100101011010210103101041010510106101071010810109101101011110112101131011410115101161011710118101191012010121101221012310124101251012610127101281012910130101311013210133101341013510136101371013810139101401014110142101431014410145101461014710148101491015010151101521015310154101551015610157101581015910160101611016210163101641016510166101671016810169101701017110172101731017410175101761017710178101791018010181101821018310184101851018610187101881018910190101911019210193101941019510196101971019810199102001020110202102031020410205102061020710208102091021010211102121021310214102151021610217102181021910220102211022210223102241022510226102271022810229102301023110232102331023410235102361023710238102391024010241102421024310244102451024610247102481024910250102511025210253102541025510256102571025810259102601026110262102631026410265102661026710268102691027010271102721027310274102751027610277102781027910280102811028210283102841028510286102871028810289102901029110292102931029410295102961029710298102991030010301103021030310304103051030610307103081030910310103111031210313103141031510316103171031810319103201032110322103231032410325103261032710328103291033010331103321033310334103351033610337103381033910340103411034210343103441034510346103471034810349103501035110352103531035410355103561035710358103591036010361103621036310364103651036610367103681036910370103711037210373103741037510376103771037810379103801038110382103831038410385103861038710388103891039010391103921039310394103951039610397103981039910400104011040210403104041040510406104071040810409104101041110412104131041410415104161041710418104191042010421104221042310424104251042610427104281042910430104311043210433104341043510436104371043810439104401044110442104431044410445104461044710448104491045010451104521045310454104551045610457104581045910460104611046210463104641046510466104671046810469104701047110472104731047410475104761047710478104791048010481104821048310484104851048610487104881048910490104911049210493104941049510496104971049810499105001050110502105031050410505105061050710508105091051010511105121051310514105151051610517105181051910520105211052210523105241052510526105271052810529105301053110532105331053410535105361053710538105391054010541105421054310544105451054610547105481054910550105511055210553105541055510556105571055810559105601056110562105631056410565105661056710568105691057010571105721057310574105751057610577105781057910580105811058210583105841058510586105871058810589105901059110592105931059410595105961059710598105991060010601106021060310604106051060610607106081060910610106111061210613106141061510616106171061810619106201062110622106231062410625106261062710628106291063010631106321063310634106351063610637106381063910640106411064210643106441064510646106471064810649106501065110652106531065410655106561065710658106591066010661106621066310664106651066610667106681066910670106711067210673106741067510676106771067810679106801068110682106831068410685106861068710688106891069010691106921069310694106951069610697106981069910700107011070210703107041070510706107071070810709107101071110712107131071410715107161071710718107191072010721107221072310724107251072610727107281072910730107311073210733107341073510736107371073810739107401074110742107431074410745107461074710748107491075010751107521075310754107551075610757107581075910760107611076210763107641076510766107671076810769107701077110772107731077410775107761077710778107791078010781107821078310784107851078610787107881078910790107911079210793107941079510796107971079810799108001080110802108031080410805108061080710808108091081010811108121081310814108151081610817108181081910820108211082210823108241082510826108271082810829108301083110832108331083410835108361083710838108391084010841108421084310844108451084610847108481084910850108511085210853108541085510856108571085810859108601086110862108631086410865108661086710868108691087010871108721087310874108751087610877108781087910880108811088210883108841088510886108871088810889108901089110892108931089410895108961089710898108991090010901109021090310904109051090610907109081090910910109111091210913109141091510916109171091810919109201092110922109231092410925109261092710928109291093010931109321093310934109351093610937109381093910940109411094210943109441094510946109471094810949109501095110952109531095410955109561095710958109591096010961109621096310964109651096610967109681096910970109711097210973109741097510976109771097810979109801098110982109831098410985109861098710988109891099010991109921099310994109951099610997109981099911000110011100211003110041100511006110071100811009110101101111012110131101411015110161101711018110191102011021110221102311024110251102611027110281102911030110311103211033110341103511036110371103811039110401104111042110431104411045110461104711048110491105011051110521105311054110551105611057110581105911060110611106211063110641106511066110671106811069110701107111072110731107411075110761107711078110791108011081110821108311084110851108611087110881108911090110911109211093110941109511096110971109811099111001110111102111031110411105111061110711108111091111011111111121111311114111151111611117111181111911120111211112211123111241112511126111271112811129111301113111132111331113411135111361113711138111391114011141111421114311144111451114611147111481114911150111511115211153111541115511156111571115811159111601116111162111631116411165111661116711168111691117011171111721117311174111751117611177111781117911180111811118211183111841118511186111871118811189111901119111192111931119411195111961119711198111991120011201112021120311204112051120611207112081120911210112111121211213112141121511216112171121811219112201122111222112231122411225112261122711228112291123011231112321123311234112351123611237112381123911240112411124211243112441124511246112471124811249112501125111252112531125411255112561125711258112591126011261112621126311264112651126611267112681126911270112711127211273112741127511276112771127811279112801128111282112831128411285112861128711288112891129011291112921129311294112951129611297112981129911300113011130211303113041130511306113071130811309113101131111312113131131411315113161131711318113191132011321113221132311324113251132611327113281132911330113311133211333113341133511336113371133811339113401134111342113431134411345113461134711348113491135011351113521135311354113551135611357113581135911360113611136211363113641136511366113671136811369113701137111372113731137411375113761137711378113791138011381113821138311384113851138611387113881138911390113911139211393113941139511396113971139811399114001140111402114031140411405114061140711408114091141011411114121141311414114151141611417114181141911420114211142211423114241142511426114271142811429114301143111432114331143411435114361143711438114391144011441114421144311444114451144611447114481144911450114511145211453114541145511456114571145811459114601146111462114631146411465114661146711468114691147011471114721147311474114751147611477114781147911480114811148211483114841148511486114871148811489114901149111492114931149411495114961149711498114991150011501115021150311504115051150611507115081150911510115111151211513115141151511516115171151811519115201152111522115231152411525115261152711528115291153011531115321153311534115351153611537115381153911540115411154211543115441154511546115471154811549115501155111552115531155411555115561155711558115591156011561115621156311564115651156611567115681156911570115711157211573115741157511576115771157811579115801158111582115831158411585115861158711588115891159011591115921159311594115951159611597115981159911600116011160211603116041160511606116071160811609116101161111612116131161411615116161161711618116191162011621116221162311624116251162611627116281162911630116311163211633116341163511636116371163811639116401164111642116431164411645116461164711648116491165011651116521165311654116551165611657116581165911660116611166211663116641166511666116671166811669116701167111672116731167411675116761167711678116791168011681116821168311684116851168611687116881168911690116911169211693116941169511696116971169811699117001170111702117031170411705117061170711708117091171011711117121171311714117151171611717117181171911720117211172211723117241172511726117271172811729117301173111732117331173411735117361173711738117391174011741117421174311744117451174611747117481174911750117511175211753117541175511756117571175811759117601176111762117631176411765117661176711768117691177011771117721177311774117751177611777117781177911780117811178211783117841178511786117871178811789117901179111792117931179411795117961179711798117991180011801118021180311804118051180611807118081180911810118111181211813118141181511816118171181811819118201182111822118231182411825118261182711828118291183011831118321183311834118351183611837118381183911840118411184211843118441184511846118471184811849118501185111852118531185411855118561185711858118591186011861118621186311864118651186611867118681186911870118711187211873118741187511876118771187811879118801188111882118831188411885118861188711888118891189011891118921189311894118951189611897118981189911900119011190211903119041190511906119071190811909119101191111912119131191411915119161191711918119191192011921119221192311924119251192611927119281192911930119311193211933119341193511936119371193811939119401194111942119431194411945119461194711948119491195011951119521195311954119551195611957119581195911960119611196211963119641196511966119671196811969119701197111972119731197411975119761197711978119791198011981119821198311984119851198611987119881198911990119911199211993119941199511996119971199811999120001200112002120031200412005120061200712008120091201012011120121201312014120151201612017120181201912020120211202212023120241202512026120271202812029120301203112032120331203412035120361203712038120391204012041120421204312044120451204612047120481204912050120511205212053120541205512056120571205812059120601206112062120631206412065120661206712068120691207012071120721207312074120751207612077120781207912080120811208212083120841208512086120871208812089120901209112092120931209412095120961209712098120991210012101121021210312104121051210612107121081210912110121111211212113121141211512116121171211812119121201212112122121231212412125121261212712128121291213012131121321213312134121351213612137121381213912140121411214212143121441214512146121471214812149121501215112152121531215412155121561215712158121591216012161121621216312164121651216612167121681216912170121711217212173121741217512176121771217812179121801218112182121831218412185121861218712188121891219012191121921219312194121951219612197121981219912200122011220212203122041220512206122071220812209122101221112212122131221412215122161221712218122191222012221122221222312224122251222612227122281222912230122311223212233122341223512236122371223812239122401224112242122431224412245122461224712248122491225012251122521225312254122551225612257122581225912260122611226212263122641226512266122671226812269122701227112272122731227412275122761227712278122791228012281122821228312284122851228612287122881228912290122911229212293122941229512296122971229812299123001230112302123031230412305123061230712308123091231012311123121231312314123151231612317123181231912320123211232212323123241232512326123271232812329123301233112332123331233412335123361233712338123391234012341123421234312344123451234612347123481234912350123511235212353123541235512356123571235812359123601236112362123631236412365123661236712368123691237012371123721237312374123751237612377123781237912380123811238212383123841238512386123871238812389123901239112392123931239412395123961239712398123991240012401124021240312404124051240612407124081240912410124111241212413124141241512416124171241812419124201242112422124231242412425124261242712428124291243012431124321243312434124351243612437124381243912440124411244212443124441244512446124471244812449124501245112452124531245412455124561245712458124591246012461124621246312464124651246612467124681246912470124711247212473124741247512476124771247812479124801248112482124831248412485124861248712488124891249012491124921249312494124951249612497124981249912500125011250212503125041250512506125071250812509125101251112512125131251412515125161251712518125191252012521125221252312524125251252612527125281252912530125311253212533125341253512536125371253812539125401254112542125431254412545125461254712548125491255012551125521255312554125551255612557125581255912560125611256212563125641256512566125671256812569125701257112572125731257412575125761257712578125791258012581125821258312584125851258612587125881258912590125911259212593125941259512596125971259812599126001260112602126031260412605126061260712608126091261012611126121261312614126151261612617126181261912620126211262212623126241262512626126271262812629126301263112632126331263412635126361263712638126391264012641126421264312644126451264612647126481264912650126511265212653126541265512656126571265812659126601266112662126631266412665126661266712668126691267012671126721267312674126751267612677126781267912680126811268212683126841268512686126871268812689126901269112692126931269412695126961269712698126991270012701127021270312704127051270612707127081270912710127111271212713127141271512716127171271812719127201272112722127231272412725127261272712728127291273012731127321273312734127351273612737127381273912740127411274212743127441274512746127471274812749127501275112752127531275412755127561275712758127591276012761127621276312764127651276612767127681276912770127711277212773127741277512776127771277812779127801278112782127831278412785127861278712788127891279012791127921279312794127951279612797127981279912800128011280212803128041280512806128071280812809128101281112812128131281412815128161281712818128191282012821128221282312824128251282612827128281282912830128311283212833128341283512836128371283812839128401284112842128431284412845128461284712848128491285012851128521285312854128551285612857128581285912860128611286212863128641286512866128671286812869128701287112872128731287412875128761287712878128791288012881128821288312884128851288612887128881288912890128911289212893128941289512896128971289812899129001290112902129031290412905129061290712908129091291012911129121291312914129151291612917129181291912920129211292212923129241292512926129271292812929129301293112932129331293412935129361293712938129391294012941129421294312944129451294612947129481294912950129511295212953129541295512956129571295812959129601296112962129631296412965129661296712968129691297012971129721297312974129751297612977129781297912980129811298212983129841298512986129871298812989129901299112992129931299412995129961299712998129991300013001130021300313004130051300613007130081300913010130111301213013130141301513016130171301813019130201302113022130231302413025130261302713028130291303013031130321303313034130351303613037130381303913040130411304213043130441304513046130471304813049130501305113052130531305413055130561305713058130591306013061130621306313064130651306613067130681306913070130711307213073130741307513076130771307813079130801308113082130831308413085130861308713088130891309013091130921309313094130951309613097130981309913100131011310213103131041310513106131071310813109131101311113112131131311413115131161311713118131191312013121131221312313124131251312613127131281312913130131311313213133131341313513136131371313813139131401314113142131431314413145131461314713148131491315013151131521315313154131551315613157131581315913160131611316213163131641316513166131671316813169131701317113172131731317413175131761317713178131791318013181131821318313184131851318613187131881318913190131911319213193131941319513196131971319813199132001320113202132031320413205132061320713208132091321013211132121321313214132151321613217132181321913220132211322213223132241322513226132271322813229132301323113232132331323413235132361323713238132391324013241132421324313244132451324613247132481324913250132511325213253132541325513256132571325813259132601326113262132631326413265132661326713268132691327013271132721327313274132751327613277132781327913280132811328213283132841328513286132871328813289132901329113292132931329413295132961329713298132991330013301133021330313304133051330613307133081330913310133111331213313133141331513316133171331813319133201332113322133231332413325133261332713328133291333013331133321333313334133351333613337133381333913340133411334213343133441334513346133471334813349133501335113352133531335413355133561335713358133591336013361133621336313364133651336613367133681336913370133711337213373133741337513376133771337813379133801338113382133831338413385133861338713388133891339013391133921339313394133951339613397133981339913400134011340213403134041340513406134071340813409134101341113412134131341413415134161341713418134191342013421134221342313424134251342613427134281342913430134311343213433134341343513436134371343813439134401344113442134431344413445134461344713448134491345013451134521345313454134551345613457134581345913460134611346213463134641346513466134671346813469134701347113472134731347413475134761347713478134791348013481134821348313484134851348613487134881348913490134911349213493134941349513496134971349813499135001350113502135031350413505135061350713508135091351013511135121351313514135151351613517135181351913520135211352213523135241352513526135271352813529135301353113532135331353413535135361353713538135391354013541135421354313544135451354613547135481354913550135511355213553135541355513556135571355813559135601356113562135631356413565135661356713568135691357013571135721357313574135751357613577135781357913580135811358213583135841358513586135871358813589135901359113592135931359413595135961359713598135991360013601136021360313604136051360613607136081360913610136111361213613136141361513616136171361813619136201362113622136231362413625136261362713628136291363013631136321363313634136351363613637136381363913640136411364213643136441364513646136471364813649136501365113652136531365413655136561365713658136591366013661136621366313664136651366613667136681366913670136711367213673136741367513676136771367813679136801368113682136831368413685136861368713688136891369013691136921369313694136951369613697136981369913700137011370213703137041370513706137071370813709137101371113712137131371413715137161371713718137191372013721137221372313724137251372613727137281372913730137311373213733137341373513736137371373813739137401374113742137431374413745137461374713748137491375013751137521375313754137551375613757137581375913760137611376213763137641376513766137671376813769137701377113772137731377413775137761377713778137791378013781137821378313784137851378613787137881378913790137911379213793137941379513796137971379813799138001380113802138031380413805138061380713808138091381013811138121381313814138151381613817138181381913820138211382213823138241382513826138271382813829138301383113832138331383413835138361383713838138391384013841138421384313844138451384613847138481384913850138511385213853138541385513856138571385813859138601386113862138631386413865138661386713868138691387013871138721387313874138751387613877138781387913880138811388213883138841388513886138871388813889138901389113892138931389413895138961389713898138991390013901139021390313904139051390613907139081390913910139111391213913139141391513916139171391813919139201392113922139231392413925139261392713928139291393013931139321393313934139351393613937139381393913940139411394213943139441394513946139471394813949139501395113952139531395413955139561395713958139591396013961139621396313964139651396613967139681396913970139711397213973139741397513976139771397813979139801398113982139831398413985139861398713988139891399013991139921399313994139951399613997139981399914000140011400214003140041400514006140071400814009140101401114012140131401414015140161401714018140191402014021140221402314024140251402614027140281402914030140311403214033140341403514036140371403814039140401404114042140431404414045140461404714048140491405014051140521405314054140551405614057140581405914060140611406214063140641406514066140671406814069140701407114072140731407414075140761407714078140791408014081140821408314084140851408614087140881408914090140911409214093140941409514096140971409814099141001410114102141031410414105141061410714108141091411014111141121411314114141151411614117141181411914120141211412214123141241412514126141271412814129141301413114132141331413414135141361413714138141391414014141141421414314144141451414614147141481414914150141511415214153141541415514156141571415814159141601416114162141631416414165141661416714168141691417014171141721417314174141751417614177141781417914180141811418214183141841418514186141871418814189141901419114192141931419414195141961419714198141991420014201142021420314204142051420614207142081420914210142111421214213142141421514216142171421814219142201422114222142231422414225142261422714228142291423014231142321423314234142351423614237142381423914240142411424214243142441424514246142471424814249142501425114252142531425414255142561425714258142591426014261142621426314264142651426614267142681426914270142711427214273142741427514276142771427814279142801428114282142831428414285142861428714288142891429014291142921429314294142951429614297142981429914300143011430214303143041430514306143071430814309143101431114312143131431414315143161431714318143191432014321143221432314324143251432614327143281432914330143311433214333143341433514336143371433814339143401434114342143431434414345143461434714348143491435014351143521435314354143551435614357143581435914360143611436214363143641436514366143671436814369143701437114372143731437414375143761437714378143791438014381143821438314384143851438614387143881438914390143911439214393143941439514396143971439814399144001440114402144031440414405144061440714408144091441014411144121441314414144151441614417144181441914420144211442214423144241442514426144271442814429144301443114432144331443414435144361443714438144391444014441144421444314444144451444614447144481444914450144511445214453144541445514456144571445814459144601446114462144631446414465144661446714468144691447014471144721447314474144751447614477144781447914480144811448214483144841448514486144871448814489144901449114492144931449414495144961449714498144991450014501145021450314504145051450614507145081450914510145111451214513145141451514516145171451814519145201452114522145231452414525145261452714528145291453014531145321453314534145351453614537145381453914540145411454214543145441454514546145471454814549145501455114552145531455414555145561455714558145591456014561145621456314564145651456614567145681456914570145711457214573145741457514576145771457814579145801458114582145831458414585145861458714588145891459014591145921459314594145951459614597145981459914600146011460214603146041460514606146071460814609146101461114612146131461414615146161461714618146191462014621146221462314624146251462614627146281462914630146311463214633146341463514636146371463814639146401464114642146431464414645146461464714648146491465014651146521465314654146551465614657146581465914660146611466214663146641466514666146671466814669146701467114672146731467414675146761467714678146791468014681146821468314684146851468614687146881468914690146911469214693146941469514696146971469814699147001470114702147031470414705147061470714708147091471014711147121471314714147151471614717147181471914720147211472214723147241472514726147271472814729147301473114732147331473414735147361473714738147391474014741147421474314744147451474614747147481474914750147511475214753147541475514756147571475814759147601476114762147631476414765147661476714768147691477014771147721477314774147751477614777147781477914780147811478214783147841478514786147871478814789147901479114792147931479414795147961479714798147991480014801148021480314804148051480614807148081480914810148111481214813148141481514816148171481814819148201482114822148231482414825148261482714828148291483014831148321483314834148351483614837148381483914840148411484214843148441484514846148471484814849148501485114852148531485414855148561485714858148591486014861148621486314864148651486614867148681486914870148711487214873148741487514876148771487814879148801488114882148831488414885148861488714888148891489014891148921489314894148951489614897148981489914900149011490214903149041490514906149071490814909149101491114912149131491414915149161491714918149191492014921149221492314924149251492614927149281492914930149311493214933149341493514936149371493814939149401494114942149431494414945149461494714948149491495014951149521495314954149551495614957149581495914960149611496214963149641496514966149671496814969149701497114972149731497414975149761497714978149791498014981149821498314984149851498614987149881498914990149911499214993149941499514996149971499814999150001500115002150031500415005150061500715008150091501015011150121501315014150151501615017150181501915020150211502215023150241502515026150271502815029150301503115032150331503415035150361503715038150391504015041150421504315044150451504615047150481504915050150511505215053150541505515056150571505815059150601506115062150631506415065150661506715068150691507015071150721507315074150751507615077150781507915080150811508215083150841508515086150871508815089150901509115092150931509415095150961509715098150991510015101151021510315104151051510615107151081510915110151111511215113151141511515116151171511815119151201512115122151231512415125151261512715128151291513015131151321513315134151351513615137151381513915140151411514215143151441514515146151471514815149151501515115152151531515415155151561515715158151591516015161151621516315164151651516615167151681516915170151711517215173151741517515176151771517815179151801518115182151831518415185151861518715188151891519015191151921519315194151951519615197151981519915200152011520215203152041520515206152071520815209152101521115212152131521415215152161521715218152191522015221152221522315224152251522615227152281522915230152311523215233152341523515236152371523815239152401524115242152431524415245152461524715248152491525015251152521525315254152551525615257152581525915260152611526215263152641526515266152671526815269152701527115272152731527415275152761527715278152791528015281152821528315284152851528615287152881528915290152911529215293152941529515296152971529815299153001530115302153031530415305153061530715308153091531015311153121531315314153151531615317153181531915320153211532215323153241532515326153271532815329153301533115332153331533415335153361533715338153391534015341153421534315344153451534615347153481534915350153511535215353153541535515356153571535815359153601536115362153631536415365153661536715368153691537015371153721537315374153751537615377153781537915380153811538215383153841538515386153871538815389153901539115392153931539415395153961539715398153991540015401154021540315404154051540615407154081540915410154111541215413154141541515416154171541815419154201542115422154231542415425154261542715428154291543015431154321543315434154351543615437154381543915440154411544215443154441544515446154471544815449154501545115452154531545415455154561545715458154591546015461154621546315464154651546615467154681546915470154711547215473154741547515476154771547815479154801548115482154831548415485154861548715488154891549015491154921549315494154951549615497154981549915500155011550215503155041550515506155071550815509155101551115512155131551415515155161551715518155191552015521155221552315524155251552615527155281552915530155311553215533155341553515536155371553815539155401554115542155431554415545155461554715548155491555015551155521555315554155551555615557155581555915560155611556215563155641556515566155671556815569155701557115572155731557415575155761557715578155791558015581155821558315584155851558615587155881558915590155911559215593155941559515596155971559815599156001560115602156031560415605156061560715608156091561015611156121561315614156151561615617156181561915620156211562215623156241562515626156271562815629156301563115632156331563415635156361563715638156391564015641156421564315644156451564615647156481564915650156511565215653156541565515656156571565815659156601566115662156631566415665156661566715668156691567015671156721567315674156751567615677156781567915680156811568215683156841568515686156871568815689156901569115692156931569415695156961569715698156991570015701157021570315704157051570615707157081570915710157111571215713157141571515716157171571815719157201572115722157231572415725157261572715728157291573015731157321573315734157351573615737157381573915740157411574215743157441574515746157471574815749157501575115752157531575415755157561575715758157591576015761157621576315764157651576615767157681576915770157711577215773157741577515776157771577815779157801578115782157831578415785157861578715788157891579015791157921579315794157951579615797157981579915800158011580215803158041580515806158071580815809158101581115812158131581415815158161581715818158191582015821158221582315824158251582615827158281582915830158311583215833158341583515836158371583815839158401584115842158431584415845158461584715848158491585015851158521585315854158551585615857158581585915860158611586215863158641586515866158671586815869158701587115872158731587415875158761587715878158791588015881158821588315884158851588615887158881588915890158911589215893158941589515896158971589815899159001590115902159031590415905159061590715908159091591015911159121591315914159151591615917159181591915920159211592215923159241592515926159271592815929159301593115932159331593415935159361593715938159391594015941159421594315944159451594615947159481594915950159511595215953159541595515956159571595815959159601596115962159631596415965159661596715968159691597015971159721597315974159751597615977159781597915980159811598215983159841598515986159871598815989159901599115992159931599415995159961599715998159991600016001160021600316004160051600616007160081600916010160111601216013160141601516016160171601816019160201602116022160231602416025160261602716028160291603016031160321603316034160351603616037160381603916040160411604216043160441604516046160471604816049160501605116052160531605416055160561605716058160591606016061160621606316064160651606616067160681606916070160711607216073160741607516076160771607816079160801608116082160831608416085160861608716088160891609016091160921609316094160951609616097160981609916100161011610216103161041610516106161071610816109161101611116112161131611416115161161611716118161191612016121161221612316124161251612616127161281612916130161311613216133161341613516136161371613816139161401614116142161431614416145161461614716148161491615016151161521615316154161551615616157161581615916160161611616216163161641616516166161671616816169161701617116172161731617416175161761617716178161791618016181161821618316184161851618616187161881618916190161911619216193161941619516196161971619816199162001620116202162031620416205162061620716208162091621016211162121621316214162151621616217162181621916220162211622216223162241622516226162271622816229162301623116232162331623416235162361623716238162391624016241162421624316244162451624616247162481624916250162511625216253162541625516256162571625816259162601626116262162631626416265162661626716268162691627016271162721627316274162751627616277162781627916280162811628216283162841628516286162871628816289162901629116292162931629416295162961629716298162991630016301163021630316304163051630616307163081630916310163111631216313163141631516316163171631816319163201632116322163231632416325163261632716328163291633016331163321633316334163351633616337163381633916340163411634216343163441634516346163471634816349163501635116352163531635416355163561635716358163591636016361163621636316364163651636616367163681636916370163711637216373163741637516376163771637816379163801638116382163831638416385163861638716388163891639016391163921639316394163951639616397163981639916400164011640216403164041640516406164071640816409164101641116412164131641416415164161641716418164191642016421164221642316424164251642616427164281642916430164311643216433164341643516436164371643816439164401644116442164431644416445164461644716448164491645016451164521645316454164551645616457164581645916460164611646216463164641646516466164671646816469164701647116472164731647416475164761647716478164791648016481164821648316484164851648616487164881648916490164911649216493164941649516496164971649816499165001650116502165031650416505165061650716508165091651016511165121651316514165151651616517165181651916520165211652216523165241652516526165271652816529165301653116532165331653416535165361653716538165391654016541165421654316544165451654616547165481654916550165511655216553165541655516556165571655816559165601656116562165631656416565165661656716568165691657016571165721657316574165751657616577165781657916580165811658216583165841658516586165871658816589165901659116592165931659416595165961659716598165991660016601166021660316604166051660616607166081660916610166111661216613166141661516616166171661816619166201662116622166231662416625166261662716628166291663016631166321663316634166351663616637166381663916640166411664216643166441664516646166471664816649166501665116652166531665416655166561665716658166591666016661166621666316664166651666616667166681666916670166711667216673166741667516676166771667816679166801668116682166831668416685166861668716688166891669016691166921669316694166951669616697166981669916700167011670216703167041670516706167071670816709167101671116712167131671416715167161671716718167191672016721167221672316724167251672616727167281672916730167311673216733167341673516736167371673816739167401674116742167431674416745167461674716748167491675016751167521675316754167551675616757167581675916760167611676216763167641676516766167671676816769167701677116772167731677416775167761677716778167791678016781167821678316784167851678616787167881678916790167911679216793167941679516796167971679816799168001680116802168031680416805168061680716808168091681016811168121681316814168151681616817168181681916820168211682216823168241682516826168271682816829168301683116832168331683416835168361683716838168391684016841168421684316844168451684616847168481684916850168511685216853168541685516856168571685816859168601686116862168631686416865168661686716868168691687016871168721687316874168751687616877168781687916880168811688216883168841688516886168871688816889168901689116892168931689416895168961689716898168991690016901169021690316904169051690616907169081690916910169111691216913169141691516916169171691816919169201692116922169231692416925169261692716928169291693016931169321693316934169351693616937169381693916940169411694216943169441694516946169471694816949169501695116952169531695416955169561695716958169591696016961169621696316964169651696616967169681696916970169711697216973169741697516976169771697816979169801698116982169831698416985169861698716988169891699016991169921699316994169951699616997169981699917000170011700217003170041700517006170071700817009170101701117012170131701417015170161701717018170191702017021170221702317024170251702617027170281702917030170311703217033170341703517036170371703817039170401704117042170431704417045170461704717048170491705017051170521705317054170551705617057170581705917060170611706217063170641706517066170671706817069170701707117072170731707417075170761707717078170791708017081170821708317084170851708617087170881708917090170911709217093170941709517096170971709817099171001710117102171031710417105171061710717108171091711017111171121711317114171151711617117171181711917120171211712217123171241712517126171271712817129171301713117132171331713417135171361713717138171391714017141171421714317144171451714617147171481714917150171511715217153171541715517156171571715817159171601716117162171631716417165171661716717168171691717017171171721717317174171751717617177171781717917180171811718217183171841718517186171871718817189171901719117192171931719417195171961719717198171991720017201172021720317204172051720617207172081720917210172111721217213172141721517216172171721817219172201722117222172231722417225172261722717228172291723017231172321723317234172351723617237172381723917240172411724217243172441724517246172471724817249172501725117252172531725417255172561725717258172591726017261172621726317264172651726617267172681726917270172711727217273172741727517276172771727817279172801728117282172831728417285172861728717288172891729017291172921729317294172951729617297172981729917300173011730217303173041730517306173071730817309173101731117312173131731417315173161731717318173191732017321173221732317324173251732617327173281732917330173311733217333173341733517336173371733817339173401734117342173431734417345173461734717348173491735017351173521735317354173551735617357173581735917360173611736217363173641736517366173671736817369173701737117372173731737417375173761737717378173791738017381173821738317384173851738617387173881738917390173911739217393173941739517396173971739817399174001740117402174031740417405174061740717408174091741017411174121741317414174151741617417174181741917420174211742217423174241742517426174271742817429174301743117432174331743417435174361743717438174391744017441174421744317444174451744617447174481744917450174511745217453174541745517456174571745817459174601746117462174631746417465174661746717468174691747017471174721747317474174751747617477174781747917480174811748217483174841748517486174871748817489174901749117492174931749417495174961749717498174991750017501175021750317504175051750617507175081750917510175111751217513175141751517516175171751817519175201752117522175231752417525175261752717528175291753017531175321753317534175351753617537175381753917540175411754217543175441754517546175471754817549175501755117552175531755417555175561755717558175591756017561175621756317564175651756617567175681756917570175711757217573175741757517576175771757817579175801758117582175831758417585175861758717588175891759017591175921759317594175951759617597175981759917600176011760217603176041760517606176071760817609176101761117612176131761417615176161761717618176191762017621176221762317624176251762617627176281762917630176311763217633176341763517636176371763817639176401764117642176431764417645176461764717648176491765017651176521765317654176551765617657176581765917660176611766217663176641766517666176671766817669176701767117672176731767417675176761767717678176791768017681176821768317684176851768617687176881768917690176911769217693176941769517696176971769817699177001770117702177031770417705177061770717708177091771017711177121771317714177151771617717177181771917720177211772217723177241772517726177271772817729177301773117732177331773417735177361773717738177391774017741177421774317744177451774617747177481774917750177511775217753177541775517756177571775817759177601776117762177631776417765177661776717768177691777017771177721777317774177751777617777177781777917780177811778217783177841778517786177871778817789177901779117792177931779417795177961779717798177991780017801178021780317804178051780617807178081780917810178111781217813178141781517816178171781817819178201782117822178231782417825178261782717828178291783017831178321783317834178351783617837178381783917840178411784217843178441784517846178471784817849178501785117852178531785417855178561785717858178591786017861178621786317864178651786617867178681786917870178711787217873178741787517876178771787817879178801788117882178831788417885178861788717888178891789017891178921789317894178951789617897178981789917900179011790217903179041790517906179071790817909179101791117912179131791417915179161791717918179191792017921179221792317924179251792617927179281792917930179311793217933179341793517936179371793817939179401794117942179431794417945179461794717948179491795017951179521795317954179551795617957179581795917960179611796217963179641796517966179671796817969179701797117972179731797417975179761797717978179791798017981179821798317984179851798617987179881798917990179911799217993179941799517996179971799817999180001800118002180031800418005180061800718008180091801018011180121801318014180151801618017180181801918020180211802218023180241802518026180271802818029180301803118032180331803418035180361803718038180391804018041180421804318044180451804618047180481804918050180511805218053180541805518056180571805818059180601806118062180631806418065180661806718068180691807018071180721807318074180751807618077180781807918080180811808218083180841808518086180871808818089180901809118092180931809418095180961809718098180991810018101181021810318104181051810618107181081810918110181111811218113181141811518116181171811818119181201812118122181231812418125181261812718128181291813018131181321813318134181351813618137181381813918140181411814218143181441814518146181471814818149181501815118152181531815418155181561815718158181591816018161181621816318164181651816618167181681816918170181711817218173181741817518176181771817818179181801818118182181831818418185181861818718188181891819018191181921819318194181951819618197181981819918200182011820218203182041820518206182071820818209182101821118212182131821418215182161821718218182191822018221182221822318224182251822618227182281822918230182311823218233182341823518236182371823818239182401824118242182431824418245182461824718248182491825018251182521825318254182551825618257182581825918260182611826218263182641826518266182671826818269182701827118272182731827418275182761827718278182791828018281182821828318284182851828618287182881828918290182911829218293182941829518296182971829818299183001830118302183031830418305183061830718308183091831018311183121831318314183151831618317183181831918320183211832218323183241832518326183271832818329183301833118332183331833418335183361833718338183391834018341183421834318344183451834618347183481834918350183511835218353183541835518356183571835818359183601836118362183631836418365183661836718368183691837018371183721837318374183751837618377183781837918380183811838218383183841838518386183871838818389183901839118392183931839418395183961839718398183991840018401184021840318404184051840618407184081840918410184111841218413184141841518416184171841818419184201842118422184231842418425184261842718428184291843018431184321843318434184351843618437184381843918440184411844218443184441844518446184471844818449184501845118452184531845418455184561845718458184591846018461184621846318464184651846618467184681846918470184711847218473184741847518476184771847818479184801848118482184831848418485184861848718488184891849018491184921849318494184951849618497184981849918500185011850218503185041850518506185071850818509185101851118512185131851418515185161851718518185191852018521185221852318524185251852618527185281852918530185311853218533185341853518536185371853818539185401854118542185431854418545185461854718548185491855018551185521855318554185551855618557185581855918560185611856218563185641856518566185671856818569185701857118572185731857418575185761857718578185791858018581185821858318584185851858618587185881858918590185911859218593185941859518596185971859818599186001860118602186031860418605186061860718608186091861018611186121861318614186151861618617186181861918620186211862218623186241862518626186271862818629186301863118632186331863418635186361863718638186391864018641186421864318644186451864618647186481864918650186511865218653186541865518656186571865818659186601866118662186631866418665186661866718668186691867018671186721867318674186751867618677186781867918680186811868218683186841868518686186871868818689186901869118692186931869418695186961869718698186991870018701187021870318704187051870618707187081870918710187111871218713187141871518716187171871818719187201872118722187231872418725187261872718728187291873018731187321873318734187351873618737187381873918740187411874218743187441874518746187471874818749187501875118752187531875418755187561875718758187591876018761187621876318764187651876618767187681876918770187711877218773187741877518776187771877818779187801878118782187831878418785187861878718788187891879018791187921879318794187951879618797187981879918800188011880218803188041880518806188071880818809188101881118812188131881418815188161881718818188191882018821188221882318824188251882618827188281882918830188311883218833188341883518836188371883818839188401884118842188431884418845188461884718848188491885018851188521885318854188551885618857188581885918860188611886218863188641886518866188671886818869188701887118872188731887418875188761887718878188791888018881188821888318884188851888618887188881888918890188911889218893188941889518896188971889818899189001890118902189031890418905189061890718908189091891018911189121891318914189151891618917189181891918920189211892218923189241892518926189271892818929189301893118932189331893418935189361893718938189391894018941189421894318944189451894618947189481894918950189511895218953189541895518956189571895818959189601896118962189631896418965189661896718968189691897018971189721897318974189751897618977189781897918980189811898218983189841898518986189871898818989189901899118992189931899418995189961899718998189991900019001190021900319004190051900619007190081900919010190111901219013190141901519016190171901819019190201902119022190231902419025190261902719028190291903019031190321903319034190351903619037190381903919040190411904219043190441904519046190471904819049190501905119052190531905419055190561905719058190591906019061190621906319064190651906619067190681906919070190711907219073190741907519076190771907819079190801908119082190831908419085190861908719088190891909019091190921909319094190951909619097190981909919100191011910219103191041910519106191071910819109191101911119112191131911419115191161911719118191191912019121191221912319124191251912619127191281912919130191311913219133191341913519136191371913819139191401914119142191431914419145191461914719148191491915019151191521915319154191551915619157191581915919160191611916219163191641916519166191671916819169191701917119172191731917419175191761917719178191791918019181191821918319184191851918619187191881918919190191911919219193191941919519196191971919819199192001920119202192031920419205192061920719208192091921019211192121921319214192151921619217192181921919220192211922219223192241922519226192271922819229192301923119232192331923419235192361923719238192391924019241192421924319244192451924619247192481924919250192511925219253192541925519256192571925819259192601926119262192631926419265192661926719268192691927019271192721927319274192751927619277192781927919280192811928219283192841928519286192871928819289192901929119292192931929419295192961929719298192991930019301193021930319304193051930619307193081930919310193111931219313193141931519316193171931819319193201932119322193231932419325193261932719328193291933019331193321933319334193351933619337193381933919340193411934219343193441934519346193471934819349193501935119352193531935419355193561935719358193591936019361193621936319364193651936619367193681936919370193711937219373193741937519376193771937819379193801938119382193831938419385193861938719388193891939019391193921939319394193951939619397193981939919400194011940219403194041940519406194071940819409194101941119412194131941419415194161941719418194191942019421194221942319424194251942619427194281942919430194311943219433194341943519436194371943819439194401944119442194431944419445194461944719448194491945019451194521945319454194551945619457194581945919460194611946219463194641946519466194671946819469194701947119472194731947419475194761947719478194791948019481194821948319484194851948619487194881948919490194911949219493194941949519496194971949819499195001950119502195031950419505195061950719508195091951019511195121951319514195151951619517195181951919520195211952219523195241952519526195271952819529195301953119532195331953419535195361953719538195391954019541195421954319544195451954619547195481954919550195511955219553195541955519556195571955819559195601956119562195631956419565195661956719568195691957019571195721957319574195751957619577195781957919580195811958219583195841958519586195871958819589195901959119592195931959419595195961959719598195991960019601196021960319604196051960619607196081960919610196111961219613196141961519616196171961819619196201962119622196231962419625196261962719628196291963019631196321963319634196351963619637196381963919640196411964219643196441964519646196471964819649196501965119652196531965419655196561965719658196591966019661196621966319664196651966619667196681966919670196711967219673196741967519676196771967819679196801968119682196831968419685196861968719688196891969019691196921969319694196951969619697196981969919700197011970219703197041970519706197071970819709197101971119712197131971419715197161971719718197191972019721197221972319724197251972619727197281972919730197311973219733197341973519736197371973819739197401974119742197431974419745197461974719748197491975019751197521975319754197551975619757197581975919760197611976219763197641976519766197671976819769197701977119772197731977419775197761977719778197791978019781197821978319784197851978619787197881978919790197911979219793197941979519796197971979819799198001980119802198031980419805198061980719808198091981019811198121981319814198151981619817198181981919820198211982219823198241982519826198271982819829198301983119832198331983419835198361983719838198391984019841198421984319844198451984619847198481984919850198511985219853198541985519856198571985819859198601986119862198631986419865198661986719868198691987019871198721987319874198751987619877198781987919880198811988219883198841988519886198871988819889198901989119892198931989419895198961989719898198991990019901199021990319904199051990619907199081990919910199111991219913199141991519916199171991819919199201992119922199231992419925199261992719928199291993019931199321993319934199351993619937199381993919940199411994219943199441994519946199471994819949199501995119952199531995419955199561995719958199591996019961199621996319964199651996619967199681996919970199711997219973199741997519976199771997819979199801998119982199831998419985199861998719988199891999019991199921999319994199951999619997199981999920000200012000220003200042000520006200072000820009200102001120012200132001420015200162001720018200192002020021200222002320024200252002620027200282002920030200312003220033200342003520036200372003820039200402004120042200432004420045200462004720048200492005020051200522005320054200552005620057200582005920060200612006220063200642006520066200672006820069200702007120072200732007420075200762007720078200792008020081200822008320084200852008620087200882008920090200912009220093200942009520096200972009820099201002010120102201032010420105201062010720108201092011020111201122011320114201152011620117201182011920120201212012220123201242012520126201272012820129201302013120132201332013420135201362013720138201392014020141201422014320144201452014620147201482014920150201512015220153201542015520156201572015820159201602016120162201632016420165201662016720168201692017020171201722017320174201752017620177201782017920180201812018220183201842018520186201872018820189201902019120192201932019420195201962019720198201992020020201202022020320204202052020620207202082020920210202112021220213202142021520216202172021820219202202022120222202232022420225202262022720228202292023020231202322023320234202352023620237202382023920240202412024220243202442024520246202472024820249202502025120252202532025420255202562025720258202592026020261202622026320264202652026620267202682026920270202712027220273202742027520276202772027820279202802028120282202832028420285202862028720288202892029020291202922029320294202952029620297202982029920300203012030220303203042030520306203072030820309203102031120312203132031420315203162031720318203192032020321203222032320324203252032620327203282032920330203312033220333203342033520336203372033820339203402034120342203432034420345203462034720348203492035020351203522035320354203552035620357203582035920360203612036220363203642036520366203672036820369203702037120372203732037420375203762037720378203792038020381203822038320384203852038620387203882038920390203912039220393203942039520396203972039820399204002040120402204032040420405204062040720408204092041020411204122041320414204152041620417204182041920420204212042220423204242042520426204272042820429204302043120432204332043420435204362043720438204392044020441204422044320444204452044620447204482044920450204512045220453204542045520456204572045820459204602046120462204632046420465204662046720468204692047020471204722047320474204752047620477204782047920480204812048220483204842048520486204872048820489204902049120492204932049420495204962049720498204992050020501205022050320504205052050620507205082050920510205112051220513205142051520516205172051820519205202052120522205232052420525205262052720528205292053020531205322053320534205352053620537205382053920540205412054220543205442054520546205472054820549205502055120552205532055420555205562055720558205592056020561205622056320564205652056620567205682056920570205712057220573205742057520576205772057820579205802058120582205832058420585205862058720588205892059020591205922059320594205952059620597205982059920600206012060220603206042060520606206072060820609206102061120612206132061420615206162061720618206192062020621206222062320624206252062620627206282062920630206312063220633206342063520636206372063820639206402064120642206432064420645206462064720648206492065020651206522065320654206552065620657206582065920660206612066220663206642066520666206672066820669206702067120672206732067420675206762067720678206792068020681206822068320684206852068620687206882068920690206912069220693206942069520696206972069820699207002070120702207032070420705207062070720708207092071020711207122071320714207152071620717207182071920720207212072220723207242072520726207272072820729207302073120732207332073420735207362073720738207392074020741207422074320744207452074620747207482074920750207512075220753207542075520756207572075820759207602076120762207632076420765207662076720768207692077020771207722077320774207752077620777207782077920780207812078220783207842078520786207872078820789207902079120792207932079420795207962079720798207992080020801208022080320804208052080620807208082080920810208112081220813208142081520816208172081820819208202082120822208232082420825208262082720828208292083020831208322083320834208352083620837208382083920840208412084220843208442084520846208472084820849208502085120852208532085420855208562085720858208592086020861208622086320864208652086620867208682086920870208712087220873208742087520876208772087820879208802088120882208832088420885208862088720888208892089020891208922089320894208952089620897208982089920900209012090220903209042090520906209072090820909209102091120912209132091420915209162091720918209192092020921209222092320924209252092620927209282092920930209312093220933209342093520936209372093820939209402094120942209432094420945209462094720948209492095020951209522095320954209552095620957209582095920960209612096220963209642096520966209672096820969209702097120972209732097420975209762097720978209792098020981209822098320984209852098620987209882098920990209912099220993209942099520996209972099820999210002100121002210032100421005210062100721008210092101021011210122101321014210152101621017210182101921020210212102221023210242102521026210272102821029210302103121032210332103421035210362103721038210392104021041210422104321044210452104621047210482104921050210512105221053210542105521056210572105821059210602106121062210632106421065210662106721068210692107021071210722107321074210752107621077210782107921080210812108221083210842108521086210872108821089210902109121092210932109421095210962109721098210992110021101211022110321104211052110621107211082110921110211112111221113211142111521116211172111821119211202112121122211232112421125211262112721128211292113021131211322113321134211352113621137211382113921140211412114221143211442114521146211472114821149211502115121152211532115421155211562115721158211592116021161211622116321164211652116621167211682116921170211712117221173211742117521176211772117821179211802118121182211832118421185211862118721188211892119021191211922119321194211952119621197211982119921200212012120221203212042120521206212072120821209212102121121212212132121421215212162121721218212192122021221212222122321224212252122621227212282122921230212312123221233212342123521236212372123821239212402124121242212432124421245212462124721248212492125021251212522125321254212552125621257212582125921260212612126221263212642126521266212672126821269212702127121272212732127421275212762127721278212792128021281212822128321284212852128621287212882128921290212912129221293212942129521296212972129821299213002130121302213032130421305213062130721308213092131021311213122131321314213152131621317213182131921320213212132221323213242132521326213272132821329213302133121332213332133421335213362133721338213392134021341213422134321344213452134621347213482134921350213512135221353213542135521356213572135821359213602136121362213632136421365213662136721368213692137021371213722137321374213752137621377213782137921380213812138221383213842138521386213872138821389213902139121392213932139421395213962139721398213992140021401214022140321404214052140621407214082140921410214112141221413214142141521416214172141821419214202142121422214232142421425214262142721428214292143021431214322143321434214352143621437214382143921440214412144221443214442144521446214472144821449214502145121452214532145421455214562145721458214592146021461214622146321464214652146621467214682146921470214712147221473214742147521476214772147821479214802148121482214832148421485214862148721488214892149021491214922149321494214952149621497214982149921500215012150221503215042150521506215072150821509215102151121512215132151421515215162151721518215192152021521215222152321524215252152621527215282152921530215312153221533215342153521536215372153821539215402154121542215432154421545215462154721548215492155021551215522155321554215552155621557215582155921560215612156221563215642156521566215672156821569215702157121572215732157421575215762157721578215792158021581215822158321584215852158621587215882158921590215912159221593215942159521596215972159821599216002160121602216032160421605216062160721608216092161021611216122161321614216152161621617216182161921620216212162221623216242162521626216272162821629216302163121632216332163421635216362163721638216392164021641216422164321644216452164621647216482164921650216512165221653216542165521656216572165821659216602166121662216632166421665216662166721668216692167021671216722167321674216752167621677216782167921680216812168221683216842168521686216872168821689216902169121692216932169421695216962169721698216992170021701217022170321704217052170621707217082170921710217112171221713217142171521716217172171821719217202172121722217232172421725217262172721728217292173021731217322173321734217352173621737217382173921740217412174221743217442174521746217472174821749217502175121752217532175421755217562175721758217592176021761217622176321764217652176621767217682176921770217712177221773217742177521776217772177821779217802178121782217832178421785217862178721788217892179021791217922179321794217952179621797217982179921800218012180221803218042180521806218072180821809218102181121812218132181421815218162181721818218192182021821218222182321824218252182621827218282182921830218312183221833218342183521836218372183821839218402184121842218432184421845218462184721848218492185021851218522185321854218552185621857218582185921860218612186221863218642186521866218672186821869218702187121872218732187421875218762187721878218792188021881218822188321884218852188621887218882188921890218912189221893218942189521896218972189821899219002190121902219032190421905219062190721908219092191021911219122191321914219152191621917219182191921920219212192221923219242192521926219272192821929219302193121932219332193421935219362193721938219392194021941219422194321944219452194621947219482194921950219512195221953219542195521956219572195821959219602196121962219632196421965219662196721968219692197021971219722197321974219752197621977219782197921980219812198221983219842198521986219872198821989219902199121992219932199421995219962199721998219992200022001220022200322004220052200622007220082200922010220112201222013220142201522016220172201822019220202202122022220232202422025220262202722028220292203022031220322203322034220352203622037220382203922040220412204222043220442204522046220472204822049220502205122052220532205422055220562205722058220592206022061220622206322064220652206622067220682206922070220712207222073220742207522076220772207822079220802208122082220832208422085220862208722088220892209022091220922209322094220952209622097220982209922100221012210222103221042210522106221072210822109221102211122112221132211422115221162211722118221192212022121221222212322124221252212622127221282212922130221312213222133221342213522136221372213822139221402214122142221432214422145221462214722148221492215022151221522215322154221552215622157221582215922160221612216222163221642216522166221672216822169221702217122172221732217422175221762217722178221792218022181221822218322184221852218622187221882218922190221912219222193221942219522196221972219822199222002220122202222032220422205222062220722208222092221022211222122221322214222152221622217222182221922220222212222222223222242222522226222272222822229222302223122232222332223422235222362223722238222392224022241222422224322244222452224622247222482224922250222512225222253222542225522256222572225822259222602226122262222632226422265222662226722268222692227022271222722227322274222752227622277222782227922280222812228222283222842228522286222872228822289222902229122292222932229422295222962229722298222992230022301223022230322304223052230622307223082230922310223112231222313223142231522316223172231822319223202232122322223232232422325223262232722328223292233022331223322233322334223352233622337223382233922340223412234222343223442234522346223472234822349223502235122352223532235422355223562235722358223592236022361223622236322364223652236622367223682236922370223712237222373223742237522376223772237822379223802238122382223832238422385223862238722388223892239022391223922239322394223952239622397223982239922400224012240222403224042240522406224072240822409224102241122412224132241422415224162241722418224192242022421224222242322424224252242622427224282242922430224312243222433224342243522436224372243822439224402244122442224432244422445224462244722448224492245022451224522245322454224552245622457224582245922460224612246222463224642246522466224672246822469224702247122472224732247422475224762247722478224792248022481224822248322484224852248622487224882248922490224912249222493224942249522496224972249822499225002250122502225032250422505225062250722508225092251022511225122251322514225152251622517225182251922520225212252222523225242252522526225272252822529225302253122532225332253422535225362253722538225392254022541225422254322544225452254622547225482254922550225512255222553225542255522556225572255822559225602256122562225632256422565225662256722568225692257022571225722257322574225752257622577225782257922580225812258222583225842258522586225872258822589225902259122592225932259422595225962259722598225992260022601226022260322604226052260622607226082260922610226112261222613226142261522616226172261822619226202262122622226232262422625226262262722628226292263022631226322263322634226352263622637226382263922640226412264222643226442264522646226472264822649226502265122652226532265422655226562265722658226592266022661226622266322664226652266622667226682266922670226712267222673226742267522676226772267822679226802268122682226832268422685226862268722688226892269022691226922269322694226952269622697226982269922700227012270222703227042270522706227072270822709227102271122712227132271422715227162271722718227192272022721227222272322724227252272622727227282272922730227312273222733227342273522736227372273822739227402274122742227432274422745227462274722748227492275022751227522275322754227552275622757227582275922760227612276222763227642276522766227672276822769227702277122772227732277422775227762277722778227792278022781227822278322784227852278622787227882278922790227912279222793227942279522796227972279822799228002280122802228032280422805228062280722808228092281022811228122281322814228152281622817228182281922820228212282222823228242282522826228272282822829228302283122832228332283422835228362283722838228392284022841228422284322844228452284622847228482284922850228512285222853228542285522856228572285822859228602286122862228632286422865228662286722868228692287022871228722287322874228752287622877228782287922880228812288222883228842288522886228872288822889228902289122892228932289422895228962289722898228992290022901229022290322904229052290622907229082290922910229112291222913229142291522916229172291822919229202292122922229232292422925229262292722928229292293022931229322293322934229352293622937229382293922940229412294222943229442294522946229472294822949229502295122952229532295422955229562295722958229592296022961229622296322964229652296622967229682296922970229712297222973229742297522976229772297822979229802298122982229832298422985229862298722988229892299022991229922299322994229952299622997229982299923000230012300223003230042300523006230072300823009230102301123012230132301423015230162301723018230192302023021230222302323024230252302623027230282302923030230312303223033230342303523036230372303823039230402304123042230432304423045230462304723048230492305023051230522305323054230552305623057230582305923060230612306223063230642306523066230672306823069230702307123072230732307423075230762307723078230792308023081230822308323084230852308623087230882308923090230912309223093230942309523096230972309823099231002310123102231032310423105231062310723108231092311023111231122311323114231152311623117231182311923120231212312223123231242312523126231272312823129231302313123132231332313423135231362313723138231392314023141231422314323144231452314623147231482314923150231512315223153231542315523156231572315823159231602316123162231632316423165231662316723168231692317023171231722317323174231752317623177231782317923180231812318223183231842318523186231872318823189231902319123192231932319423195231962319723198231992320023201232022320323204232052320623207232082320923210232112321223213232142321523216232172321823219232202322123222232232322423225232262322723228232292323023231232322323323234232352323623237232382323923240232412324223243232442324523246232472324823249232502325123252232532325423255232562325723258232592326023261232622326323264232652326623267232682326923270232712327223273232742327523276232772327823279232802328123282232832328423285232862328723288232892329023291232922329323294232952329623297232982329923300233012330223303233042330523306233072330823309233102331123312233132331423315233162331723318233192332023321233222332323324233252332623327233282332923330233312333223333233342333523336233372333823339233402334123342233432334423345233462334723348233492335023351233522335323354233552335623357233582335923360233612336223363233642336523366233672336823369233702337123372233732337423375233762337723378233792338023381233822338323384233852338623387233882338923390233912339223393233942339523396233972339823399234002340123402234032340423405234062340723408234092341023411234122341323414234152341623417234182341923420234212342223423234242342523426234272342823429234302343123432234332343423435234362343723438234392344023441234422344323444234452344623447234482344923450234512345223453234542345523456234572345823459234602346123462234632346423465234662346723468234692347023471234722347323474234752347623477234782347923480234812348223483234842348523486234872348823489234902349123492234932349423495234962349723498234992350023501235022350323504235052350623507235082350923510235112351223513235142351523516235172351823519235202352123522235232352423525235262352723528235292353023531235322353323534235352353623537235382353923540235412354223543235442354523546235472354823549235502355123552235532355423555235562355723558235592356023561235622356323564235652356623567235682356923570235712357223573235742357523576235772357823579235802358123582235832358423585235862358723588235892359023591235922359323594235952359623597235982359923600236012360223603236042360523606236072360823609236102361123612236132361423615236162361723618236192362023621236222362323624236252362623627236282362923630236312363223633236342363523636236372363823639236402364123642236432364423645236462364723648236492365023651236522365323654236552365623657236582365923660236612366223663236642366523666236672366823669236702367123672236732367423675236762367723678236792368023681236822368323684236852368623687236882368923690236912369223693236942369523696236972369823699237002370123702237032370423705237062370723708237092371023711237122371323714237152371623717237182371923720237212372223723237242372523726237272372823729237302373123732237332373423735237362373723738237392374023741237422374323744237452374623747237482374923750237512375223753237542375523756237572375823759237602376123762237632376423765237662376723768237692377023771237722377323774237752377623777237782377923780237812378223783237842378523786237872378823789237902379123792237932379423795237962379723798237992380023801238022380323804238052380623807238082380923810238112381223813238142381523816238172381823819238202382123822238232382423825238262382723828238292383023831238322383323834238352383623837238382383923840238412384223843238442384523846238472384823849238502385123852238532385423855238562385723858238592386023861238622386323864238652386623867238682386923870238712387223873238742387523876238772387823879238802388123882238832388423885238862388723888238892389023891238922389323894238952389623897238982389923900239012390223903239042390523906239072390823909239102391123912239132391423915239162391723918239192392023921239222392323924239252392623927239282392923930239312393223933239342393523936239372393823939239402394123942239432394423945239462394723948239492395023951239522395323954239552395623957239582395923960239612396223963239642396523966239672396823969239702397123972239732397423975239762397723978239792398023981239822398323984239852398623987239882398923990239912399223993239942399523996239972399823999240002400124002240032400424005240062400724008240092401024011240122401324014240152401624017240182401924020240212402224023240242402524026240272402824029240302403124032240332403424035240362403724038240392404024041240422404324044240452404624047240482404924050240512405224053240542405524056240572405824059240602406124062240632406424065240662406724068240692407024071240722407324074240752407624077240782407924080240812408224083240842408524086240872408824089240902409124092240932409424095240962409724098240992410024101241022410324104241052410624107241082410924110241112411224113241142411524116241172411824119241202412124122241232412424125241262412724128241292413024131241322413324134241352413624137241382413924140241412414224143241442414524146241472414824149241502415124152241532415424155241562415724158241592416024161241622416324164241652416624167241682416924170241712417224173241742417524176241772417824179241802418124182241832418424185241862418724188241892419024191241922419324194241952419624197241982419924200242012420224203242042420524206242072420824209242102421124212242132421424215242162421724218242192422024221242222422324224242252422624227242282422924230242312423224233242342423524236242372423824239242402424124242242432424424245242462424724248242492425024251242522425324254242552425624257242582425924260242612426224263242642426524266242672426824269242702427124272242732427424275242762427724278242792428024281242822428324284242852428624287242882428924290242912429224293242942429524296242972429824299243002430124302243032430424305243062430724308243092431024311243122431324314243152431624317243182431924320243212432224323243242432524326243272432824329243302433124332243332433424335243362433724338243392434024341243422434324344243452434624347243482434924350243512435224353243542435524356243572435824359243602436124362243632436424365243662436724368243692437024371243722437324374243752437624377243782437924380243812438224383243842438524386243872438824389243902439124392243932439424395243962439724398243992440024401244022440324404244052440624407244082440924410244112441224413244142441524416244172441824419244202442124422244232442424425244262442724428244292443024431244322443324434244352443624437244382443924440244412444224443244442444524446244472444824449244502445124452244532445424455244562445724458244592446024461244622446324464244652446624467244682446924470244712447224473244742447524476244772447824479244802448124482244832448424485244862448724488244892449024491244922449324494244952449624497244982449924500245012450224503245042450524506245072450824509245102451124512245132451424515245162451724518245192452024521245222452324524245252452624527245282452924530245312453224533245342453524536245372453824539245402454124542245432454424545245462454724548245492455024551245522455324554245552455624557245582455924560245612456224563245642456524566245672456824569245702457124572245732457424575245762457724578245792458024581245822458324584245852458624587245882458924590245912459224593245942459524596245972459824599246002460124602246032460424605246062460724608246092461024611246122461324614246152461624617246182461924620246212462224623246242462524626246272462824629246302463124632246332463424635246362463724638246392464024641246422464324644246452464624647246482464924650246512465224653246542465524656246572465824659246602466124662246632466424665246662466724668246692467024671246722467324674246752467624677246782467924680246812468224683246842468524686246872468824689246902469124692246932469424695246962469724698246992470024701247022470324704247052470624707247082470924710247112471224713247142471524716247172471824719247202472124722247232472424725247262472724728247292473024731247322473324734247352473624737247382473924740247412474224743247442474524746247472474824749247502475124752247532475424755247562475724758247592476024761247622476324764247652476624767247682476924770247712477224773247742477524776247772477824779247802478124782247832478424785247862478724788247892479024791247922479324794247952479624797247982479924800248012480224803248042480524806248072480824809248102481124812248132481424815248162481724818248192482024821248222482324824248252482624827248282482924830248312483224833248342483524836248372483824839248402484124842248432484424845248462484724848248492485024851248522485324854248552485624857248582485924860248612486224863248642486524866248672486824869248702487124872248732487424875248762487724878248792488024881248822488324884248852488624887248882488924890248912489224893248942489524896248972489824899249002490124902249032490424905249062490724908249092491024911249122491324914249152491624917249182491924920249212492224923249242492524926249272492824929249302493124932249332493424935249362493724938249392494024941249422494324944249452494624947249482494924950249512495224953249542495524956249572495824959249602496124962249632496424965249662496724968249692497024971249722497324974249752497624977249782497924980249812498224983249842498524986249872498824989249902499124992249932499424995249962499724998249992500025001250022500325004250052500625007250082500925010250112501225013250142501525016250172501825019250202502125022250232502425025250262502725028250292503025031250322503325034250352503625037250382503925040250412504225043250442504525046250472504825049250502505125052250532505425055250562505725058250592506025061250622506325064250652506625067250682506925070250712507225073250742507525076250772507825079250802508125082250832508425085250862508725088250892509025091250922509325094250952509625097250982509925100251012510225103251042510525106251072510825109251102511125112251132511425115251162511725118251192512025121251222512325124251252512625127251282512925130251312513225133251342513525136251372513825139251402514125142251432514425145251462514725148251492515025151251522515325154251552515625157251582515925160251612516225163251642516525166251672516825169251702517125172251732517425175251762517725178251792518025181251822518325184251852518625187251882518925190251912519225193251942519525196251972519825199252002520125202252032520425205252062520725208252092521025211252122521325214252152521625217252182521925220252212522225223252242522525226252272522825229252302523125232252332523425235252362523725238252392524025241252422524325244252452524625247252482524925250252512525225253252542525525256252572525825259252602526125262252632526425265252662526725268252692527025271252722527325274252752527625277252782527925280252812528225283252842528525286252872528825289252902529125292252932529425295252962529725298252992530025301253022530325304253052530625307253082530925310253112531225313253142531525316253172531825319253202532125322253232532425325253262532725328253292533025331253322533325334253352533625337253382533925340253412534225343253442534525346253472534825349253502535125352253532535425355253562535725358253592536025361253622536325364253652536625367253682536925370253712537225373253742537525376253772537825379253802538125382253832538425385253862538725388253892539025391253922539325394253952539625397253982539925400254012540225403254042540525406254072540825409254102541125412254132541425415254162541725418254192542025421254222542325424254252542625427254282542925430254312543225433254342543525436254372543825439254402544125442254432544425445254462544725448254492545025451254522545325454254552545625457254582545925460254612546225463254642546525466254672546825469254702547125472254732547425475254762547725478254792548025481254822548325484254852548625487254882548925490254912549225493254942549525496254972549825499255002550125502255032550425505255062550725508255092551025511255122551325514255152551625517255182551925520255212552225523255242552525526255272552825529255302553125532255332553425535255362553725538255392554025541255422554325544255452554625547255482554925550255512555225553255542555525556255572555825559255602556125562255632556425565255662556725568255692557025571255722557325574255752557625577255782557925580255812558225583255842558525586255872558825589255902559125592255932559425595255962559725598255992560025601256022560325604256052560625607256082560925610256112561225613256142561525616256172561825619256202562125622256232562425625256262562725628256292563025631256322563325634256352563625637256382563925640256412564225643256442564525646256472564825649256502565125652256532565425655256562565725658256592566025661256622566325664256652566625667256682566925670256712567225673256742567525676256772567825679256802568125682256832568425685256862568725688256892569025691256922569325694256952569625697256982569925700257012570225703257042570525706257072570825709257102571125712257132571425715257162571725718257192572025721257222572325724257252572625727257282572925730257312573225733257342573525736257372573825739257402574125742257432574425745257462574725748257492575025751257522575325754257552575625757257582575925760257612576225763257642576525766257672576825769257702577125772257732577425775257762577725778257792578025781257822578325784257852578625787257882578925790257912579225793257942579525796257972579825799258002580125802258032580425805258062580725808258092581025811258122581325814258152581625817258182581925820258212582225823258242582525826258272582825829258302583125832258332583425835258362583725838258392584025841258422584325844258452584625847258482584925850258512585225853258542585525856258572585825859258602586125862258632586425865258662586725868258692587025871258722587325874258752587625877258782587925880258812588225883258842588525886258872588825889258902589125892258932589425895258962589725898258992590025901259022590325904259052590625907259082590925910259112591225913259142591525916259172591825919259202592125922259232592425925259262592725928259292593025931259322593325934259352593625937259382593925940259412594225943259442594525946259472594825949259502595125952259532595425955259562595725958259592596025961259622596325964259652596625967259682596925970259712597225973259742597525976259772597825979259802598125982259832598425985259862598725988259892599025991259922599325994259952599625997259982599926000260012600226003260042600526006260072600826009260102601126012260132601426015260162601726018260192602026021260222602326024260252602626027260282602926030260312603226033260342603526036260372603826039260402604126042260432604426045260462604726048260492605026051260522605326054260552605626057260582605926060260612606226063260642606526066260672606826069260702607126072260732607426075260762607726078260792608026081260822608326084260852608626087260882608926090260912609226093260942609526096260972609826099261002610126102261032610426105261062610726108261092611026111261122611326114261152611626117261182611926120261212612226123261242612526126261272612826129261302613126132261332613426135261362613726138261392614026141261422614326144261452614626147261482614926150261512615226153261542615526156261572615826159261602616126162261632616426165261662616726168261692617026171261722617326174261752617626177261782617926180261812618226183261842618526186261872618826189261902619126192261932619426195261962619726198261992620026201262022620326204262052620626207262082620926210262112621226213262142621526216262172621826219262202622126222262232622426225262262622726228262292623026231262322623326234262352623626237262382623926240262412624226243262442624526246262472624826249262502625126252262532625426255262562625726258262592626026261262622626326264262652626626267262682626926270262712627226273262742627526276262772627826279262802628126282262832628426285262862628726288262892629026291262922629326294262952629626297262982629926300263012630226303263042630526306263072630826309263102631126312263132631426315263162631726318263192632026321263222632326324263252632626327263282632926330263312633226333263342633526336263372633826339263402634126342263432634426345263462634726348263492635026351263522635326354263552635626357263582635926360263612636226363263642636526366263672636826369263702637126372263732637426375263762637726378263792638026381263822638326384263852638626387263882638926390263912639226393263942639526396263972639826399264002640126402264032640426405264062640726408264092641026411264122641326414264152641626417264182641926420264212642226423264242642526426264272642826429264302643126432264332643426435264362643726438264392644026441264422644326444264452644626447264482644926450264512645226453264542645526456264572645826459264602646126462264632646426465264662646726468264692647026471264722647326474264752647626477264782647926480264812648226483264842648526486264872648826489264902649126492264932649426495264962649726498264992650026501265022650326504265052650626507265082650926510265112651226513265142651526516265172651826519265202652126522265232652426525265262652726528265292653026531265322653326534265352653626537265382653926540265412654226543265442654526546265472654826549265502655126552265532655426555265562655726558265592656026561265622656326564265652656626567265682656926570265712657226573265742657526576265772657826579265802658126582265832658426585265862658726588265892659026591265922659326594265952659626597265982659926600266012660226603266042660526606266072660826609266102661126612266132661426615266162661726618266192662026621266222662326624266252662626627266282662926630266312663226633266342663526636266372663826639266402664126642266432664426645266462664726648266492665026651266522665326654266552665626657266582665926660266612666226663266642666526666266672666826669266702667126672266732667426675266762667726678266792668026681266822668326684266852668626687266882668926690266912669226693266942669526696266972669826699267002670126702267032670426705267062670726708267092671026711267122671326714267152671626717267182671926720267212672226723267242672526726267272672826729267302673126732267332673426735267362673726738267392674026741267422674326744267452674626747267482674926750267512675226753267542675526756267572675826759267602676126762267632676426765267662676726768267692677026771267722677326774267752677626777267782677926780267812678226783267842678526786267872678826789267902679126792267932679426795267962679726798267992680026801268022680326804268052680626807268082680926810268112681226813268142681526816268172681826819268202682126822268232682426825268262682726828268292683026831268322683326834268352683626837268382683926840268412684226843268442684526846268472684826849268502685126852268532685426855268562685726858268592686026861268622686326864268652686626867268682686926870268712687226873268742687526876268772687826879268802688126882268832688426885268862688726888268892689026891268922689326894268952689626897268982689926900269012690226903269042690526906269072690826909269102691126912269132691426915269162691726918269192692026921269222692326924269252692626927269282692926930269312693226933269342693526936269372693826939269402694126942269432694426945269462694726948269492695026951269522695326954269552695626957269582695926960269612696226963269642696526966269672696826969269702697126972269732697426975269762697726978269792698026981269822698326984269852698626987269882698926990269912699226993269942699526996269972699826999270002700127002270032700427005270062700727008270092701027011270122701327014270152701627017270182701927020270212702227023270242702527026270272702827029270302703127032270332703427035270362703727038270392704027041270422704327044270452704627047270482704927050270512705227053270542705527056270572705827059270602706127062270632706427065270662706727068270692707027071270722707327074270752707627077270782707927080270812708227083270842708527086270872708827089270902709127092270932709427095270962709727098270992710027101271022710327104271052710627107271082710927110271112711227113271142711527116271172711827119271202712127122271232712427125271262712727128271292713027131271322713327134271352713627137271382713927140271412714227143271442714527146271472714827149271502715127152271532715427155271562715727158271592716027161271622716327164271652716627167271682716927170271712717227173271742717527176271772717827179271802718127182271832718427185271862718727188271892719027191271922719327194271952719627197271982719927200272012720227203272042720527206272072720827209272102721127212272132721427215272162721727218272192722027221272222722327224272252722627227272282722927230272312723227233272342723527236272372723827239272402724127242272432724427245272462724727248272492725027251272522725327254272552725627257272582725927260272612726227263272642726527266272672726827269272702727127272272732727427275272762727727278272792728027281272822728327284272852728627287272882728927290272912729227293272942729527296272972729827299273002730127302273032730427305273062730727308273092731027311273122731327314273152731627317273182731927320273212732227323273242732527326273272732827329273302733127332273332733427335273362733727338273392734027341273422734327344273452734627347273482734927350273512735227353273542735527356273572735827359273602736127362273632736427365273662736727368273692737027371273722737327374273752737627377273782737927380273812738227383273842738527386273872738827389273902739127392273932739427395273962739727398273992740027401274022740327404274052740627407274082740927410274112741227413274142741527416274172741827419274202742127422274232742427425274262742727428274292743027431274322743327434274352743627437274382743927440274412744227443274442744527446274472744827449274502745127452274532745427455274562745727458274592746027461274622746327464274652746627467274682746927470274712747227473274742747527476274772747827479274802748127482274832748427485274862748727488274892749027491274922749327494274952749627497274982749927500275012750227503275042750527506275072750827509275102751127512275132751427515275162751727518275192752027521275222752327524275252752627527275282752927530275312753227533275342753527536275372753827539275402754127542275432754427545275462754727548275492755027551275522755327554275552755627557275582755927560275612756227563275642756527566275672756827569275702757127572275732757427575275762757727578275792758027581275822758327584275852758627587275882758927590275912759227593275942759527596275972759827599276002760127602276032760427605276062760727608276092761027611276122761327614276152761627617276182761927620276212762227623276242762527626276272762827629276302763127632276332763427635276362763727638276392764027641276422764327644276452764627647276482764927650276512765227653276542765527656276572765827659276602766127662276632766427665276662766727668276692767027671276722767327674276752767627677276782767927680276812768227683276842768527686276872768827689276902769127692276932769427695276962769727698276992770027701277022770327704277052770627707277082770927710277112771227713277142771527716277172771827719277202772127722277232772427725277262772727728277292773027731277322773327734277352773627737277382773927740277412774227743277442774527746277472774827749277502775127752277532775427755277562775727758277592776027761277622776327764277652776627767277682776927770277712777227773277742777527776277772777827779277802778127782277832778427785277862778727788277892779027791277922779327794277952779627797277982779927800278012780227803278042780527806278072780827809278102781127812278132781427815278162781727818278192782027821278222782327824278252782627827278282782927830278312783227833278342783527836278372783827839278402784127842278432784427845278462784727848278492785027851278522785327854278552785627857278582785927860278612786227863278642786527866278672786827869278702787127872278732787427875278762787727878278792788027881278822788327884278852788627887278882788927890278912789227893278942789527896278972789827899279002790127902279032790427905279062790727908279092791027911279122791327914279152791627917279182791927920279212792227923279242792527926279272792827929279302793127932279332793427935279362793727938279392794027941279422794327944279452794627947279482794927950279512795227953279542795527956279572795827959279602796127962279632796427965279662796727968279692797027971279722797327974279752797627977279782797927980279812798227983279842798527986279872798827989279902799127992279932799427995279962799727998279992800028001280022800328004280052800628007280082800928010280112801228013280142801528016280172801828019280202802128022280232802428025280262802728028280292803028031280322803328034280352803628037280382803928040280412804228043280442804528046280472804828049280502805128052280532805428055280562805728058280592806028061280622806328064280652806628067280682806928070280712807228073280742807528076280772807828079280802808128082280832808428085280862808728088280892809028091280922809328094280952809628097280982809928100281012810228103281042810528106281072810828109281102811128112281132811428115281162811728118281192812028121281222812328124281252812628127281282812928130281312813228133281342813528136281372813828139281402814128142281432814428145281462814728148281492815028151281522815328154281552815628157
  1. #pragma once
  2. #include <ATen/Operators.h>
  3. #include <ATen/functorch/PlumbingHelper.h>
  4. namespace at { namespace functorch {
  5. template <typename batch_rule_t, batch_rule_t batch_rule>
  6. at::Tensor _cast_Byte_generated_plumbing(const at::Tensor & self, bool non_blocking) {
  7. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8. auto maybe_layer = maybeCurrentDynamicLayer();
  9. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10. int64_t cur_level = maybe_layer->layerId();
  11. if (!isBatchedAtLevel(self, cur_level)) {
  12. return at::_ops::_cast_Byte::call(self, non_blocking);
  13. }
  14. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15. auto results = batch_rule(self_value, self_bdim, non_blocking);
  16. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17. }
  18. template <typename batch_rule_t, batch_rule_t batch_rule>
  19. at::Tensor _cast_Char_generated_plumbing(const at::Tensor & self, bool non_blocking) {
  20. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21. auto maybe_layer = maybeCurrentDynamicLayer();
  22. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23. int64_t cur_level = maybe_layer->layerId();
  24. if (!isBatchedAtLevel(self, cur_level)) {
  25. return at::_ops::_cast_Char::call(self, non_blocking);
  26. }
  27. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  28. auto results = batch_rule(self_value, self_bdim, non_blocking);
  29. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  30. }
  31. template <typename batch_rule_t, batch_rule_t batch_rule>
  32. at::Tensor _cast_Double_generated_plumbing(const at::Tensor & self, bool non_blocking) {
  33. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  34. auto maybe_layer = maybeCurrentDynamicLayer();
  35. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  36. int64_t cur_level = maybe_layer->layerId();
  37. if (!isBatchedAtLevel(self, cur_level)) {
  38. return at::_ops::_cast_Double::call(self, non_blocking);
  39. }
  40. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  41. auto results = batch_rule(self_value, self_bdim, non_blocking);
  42. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  43. }
  44. template <typename batch_rule_t, batch_rule_t batch_rule>
  45. at::Tensor _cast_Float_generated_plumbing(const at::Tensor & self, bool non_blocking) {
  46. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  47. auto maybe_layer = maybeCurrentDynamicLayer();
  48. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  49. int64_t cur_level = maybe_layer->layerId();
  50. if (!isBatchedAtLevel(self, cur_level)) {
  51. return at::_ops::_cast_Float::call(self, non_blocking);
  52. }
  53. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  54. auto results = batch_rule(self_value, self_bdim, non_blocking);
  55. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  56. }
  57. template <typename batch_rule_t, batch_rule_t batch_rule>
  58. at::Tensor _cast_Int_generated_plumbing(const at::Tensor & self, bool non_blocking) {
  59. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  60. auto maybe_layer = maybeCurrentDynamicLayer();
  61. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  62. int64_t cur_level = maybe_layer->layerId();
  63. if (!isBatchedAtLevel(self, cur_level)) {
  64. return at::_ops::_cast_Int::call(self, non_blocking);
  65. }
  66. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  67. auto results = batch_rule(self_value, self_bdim, non_blocking);
  68. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  69. }
  70. template <typename batch_rule_t, batch_rule_t batch_rule>
  71. at::Tensor _cast_Long_generated_plumbing(const at::Tensor & self, bool non_blocking) {
  72. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  73. auto maybe_layer = maybeCurrentDynamicLayer();
  74. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  75. int64_t cur_level = maybe_layer->layerId();
  76. if (!isBatchedAtLevel(self, cur_level)) {
  77. return at::_ops::_cast_Long::call(self, non_blocking);
  78. }
  79. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  80. auto results = batch_rule(self_value, self_bdim, non_blocking);
  81. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  82. }
  83. template <typename batch_rule_t, batch_rule_t batch_rule>
  84. at::Tensor _cast_Short_generated_plumbing(const at::Tensor & self, bool non_blocking) {
  85. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  86. auto maybe_layer = maybeCurrentDynamicLayer();
  87. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  88. int64_t cur_level = maybe_layer->layerId();
  89. if (!isBatchedAtLevel(self, cur_level)) {
  90. return at::_ops::_cast_Short::call(self, non_blocking);
  91. }
  92. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  93. auto results = batch_rule(self_value, self_bdim, non_blocking);
  94. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  95. }
  96. template <typename batch_rule_t, batch_rule_t batch_rule>
  97. at::Tensor _cast_Half_generated_plumbing(const at::Tensor & self, bool non_blocking) {
  98. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  99. auto maybe_layer = maybeCurrentDynamicLayer();
  100. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  101. int64_t cur_level = maybe_layer->layerId();
  102. if (!isBatchedAtLevel(self, cur_level)) {
  103. return at::_ops::_cast_Half::call(self, non_blocking);
  104. }
  105. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  106. auto results = batch_rule(self_value, self_bdim, non_blocking);
  107. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  108. }
  109. template <typename batch_rule_t, batch_rule_t batch_rule>
  110. void _backward_generated_plumbing(const at::Tensor & self, at::TensorList inputs, const ::std::optional<at::Tensor> & gradient, ::std::optional<bool> retain_graph, bool create_graph) {
  111. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  112. auto maybe_layer = maybeCurrentDynamicLayer();
  113. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  114. int64_t cur_level = maybe_layer->layerId();
  115. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(inputs, cur_level) && !isBatchedAtLevel(gradient, cur_level)) {
  116. return at::_ops::_backward::call(self, inputs, gradient, retain_graph, create_graph);
  117. }
  118. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  119. std::optional<Tensor> gradient_value;
  120. std::optional<int64_t> gradient_bdim;
  121. if (gradient) {
  122. std::tie(gradient_value, gradient_bdim) = unwrapTensorAtLevel(gradient.value(), cur_level);
  123. }
  124. batch_rule(self_value, self_bdim, inputs, gradient_value, gradient_bdim, retain_graph, create_graph);
  125. }
  126. template <typename batch_rule_t, batch_rule_t batch_rule>
  127. void set_data_generated_plumbing(at::Tensor & self, const at::Tensor & new_data) {
  128. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  129. auto maybe_layer = maybeCurrentDynamicLayer();
  130. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  131. int64_t cur_level = maybe_layer->layerId();
  132. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(new_data, cur_level)) {
  133. return at::_ops::set_data::call(self, new_data);
  134. }
  135. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  136. auto [new_data_value, new_data_bdim] = unwrapTensorAtLevel(new_data, cur_level);
  137. batch_rule(self_value, self_bdim, new_data_value, new_data_bdim);
  138. }
  139. template <typename batch_rule_t, batch_rule_t batch_rule>
  140. at::Tensor data_generated_plumbing(const at::Tensor & self) {
  141. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  142. auto maybe_layer = maybeCurrentDynamicLayer();
  143. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  144. int64_t cur_level = maybe_layer->layerId();
  145. if (!isBatchedAtLevel(self, cur_level)) {
  146. return at::_ops::data::call(self);
  147. }
  148. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  149. auto results = batch_rule(self_value, self_bdim);
  150. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  151. }
  152. template <typename batch_rule_t, batch_rule_t batch_rule>
  153. at::Tensor & requires_grad__generated_plumbing(at::Tensor & self, bool requires_grad) {
  154. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  155. auto maybe_layer = maybeCurrentDynamicLayer();
  156. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  157. int64_t cur_level = maybe_layer->layerId();
  158. if (!isBatchedAtLevel(self, cur_level)) {
  159. return at::_ops::requires_grad_::call(self, requires_grad);
  160. }
  161. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  162. batch_rule(self_value, self_bdim, requires_grad);
  163. return self;
  164. }
  165. template <typename batch_rule_t, batch_rule_t batch_rule>
  166. void retain_grad_generated_plumbing(at::Tensor & self) {
  167. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  168. auto maybe_layer = maybeCurrentDynamicLayer();
  169. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  170. int64_t cur_level = maybe_layer->layerId();
  171. if (!isBatchedAtLevel(self, cur_level)) {
  172. return at::_ops::retain_grad::call(self);
  173. }
  174. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  175. batch_rule(self_value, self_bdim);
  176. }
  177. template <typename batch_rule_t, batch_rule_t batch_rule>
  178. at::Tensor _fw_primal_generated_plumbing(const at::Tensor & self, int64_t level) {
  179. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  180. auto maybe_layer = maybeCurrentDynamicLayer();
  181. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  182. int64_t cur_level = maybe_layer->layerId();
  183. if (!isBatchedAtLevel(self, cur_level)) {
  184. return at::_ops::_fw_primal::call(self, level);
  185. }
  186. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  187. auto results = batch_rule(self_value, self_bdim, level);
  188. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  189. }
  190. template <typename batch_rule_t, batch_rule_t batch_rule>
  191. at::Tensor _make_dual_generated_plumbing(const at::Tensor & primal, const at::Tensor & tangent, int64_t level) {
  192. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  193. auto maybe_layer = maybeCurrentDynamicLayer();
  194. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  195. int64_t cur_level = maybe_layer->layerId();
  196. if (!isBatchedAtLevel(primal, cur_level) && !isBatchedAtLevel(tangent, cur_level)) {
  197. return at::_ops::_make_dual::call(primal, tangent, level);
  198. }
  199. auto [primal_value, primal_bdim] = unwrapTensorAtLevel(primal, cur_level);
  200. auto [tangent_value, tangent_bdim] = unwrapTensorAtLevel(tangent, cur_level);
  201. auto results = batch_rule(primal_value, primal_bdim, tangent_value, tangent_bdim, level);
  202. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  203. }
  204. template <typename batch_rule_t, batch_rule_t batch_rule>
  205. ::std::tuple<at::Tensor,at::Tensor> _unpack_dual_generated_plumbing(const at::Tensor & dual, int64_t level) {
  206. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  207. auto maybe_layer = maybeCurrentDynamicLayer();
  208. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  209. int64_t cur_level = maybe_layer->layerId();
  210. if (!isBatchedAtLevel(dual, cur_level)) {
  211. return at::_ops::_unpack_dual::call(dual, level);
  212. }
  213. auto [dual_value, dual_bdim] = unwrapTensorAtLevel(dual, cur_level);
  214. auto results = batch_rule(dual_value, dual_bdim, level);
  215. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  216. }
  217. template <typename batch_rule_t, batch_rule_t batch_rule>
  218. at::Tensor _new_zeros_with_same_feature_meta_generated_plumbing(const at::Tensor & self, const at::Tensor & other, int64_t self_num_batch_dims) {
  219. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  220. auto maybe_layer = maybeCurrentDynamicLayer();
  221. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  222. int64_t cur_level = maybe_layer->layerId();
  223. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  224. return at::_ops::_new_zeros_with_same_feature_meta::call(self, other, self_num_batch_dims);
  225. }
  226. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  227. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  228. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim, self_num_batch_dims);
  229. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  230. }
  231. template <typename batch_rule_t, batch_rule_t batch_rule>
  232. at::Tensor rename_generated_plumbing(const at::Tensor & self, ::std::optional<at::DimnameList> names) {
  233. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  234. auto maybe_layer = maybeCurrentDynamicLayer();
  235. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  236. int64_t cur_level = maybe_layer->layerId();
  237. if (!isBatchedAtLevel(self, cur_level)) {
  238. return at::_ops::rename::call(self, names);
  239. }
  240. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  241. auto results = batch_rule(self_value, self_bdim, names);
  242. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  243. }
  244. template <typename batch_rule_t, batch_rule_t batch_rule>
  245. at::Tensor align_to_generated_plumbing(const at::Tensor & self, at::DimnameList names) {
  246. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  247. auto maybe_layer = maybeCurrentDynamicLayer();
  248. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  249. int64_t cur_level = maybe_layer->layerId();
  250. if (!isBatchedAtLevel(self, cur_level)) {
  251. return at::_ops::align_to::call(self, names);
  252. }
  253. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  254. auto results = batch_rule(self_value, self_bdim, names);
  255. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  256. }
  257. template <typename batch_rule_t, batch_rule_t batch_rule>
  258. at::Tensor align_to_ellipsis_idx_generated_plumbing(const at::Tensor & self, at::DimnameList order, int64_t ellipsis_idx) {
  259. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  260. auto maybe_layer = maybeCurrentDynamicLayer();
  261. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  262. int64_t cur_level = maybe_layer->layerId();
  263. if (!isBatchedAtLevel(self, cur_level)) {
  264. return at::_ops::align_to_ellipsis_idx::call(self, order, ellipsis_idx);
  265. }
  266. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  267. auto results = batch_rule(self_value, self_bdim, order, ellipsis_idx);
  268. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  269. }
  270. template <typename batch_rule_t, batch_rule_t batch_rule>
  271. at::Tensor align_as_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  272. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  273. auto maybe_layer = maybeCurrentDynamicLayer();
  274. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  275. int64_t cur_level = maybe_layer->layerId();
  276. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  277. return at::_ops::align_as::call(self, other);
  278. }
  279. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  280. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  281. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  282. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  283. }
  284. template <typename batch_rule_t, batch_rule_t batch_rule>
  285. ::std::vector<at::Tensor> align_tensors_generated_plumbing(at::TensorList tensors) {
  286. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  287. auto maybe_layer = maybeCurrentDynamicLayer();
  288. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  289. int64_t cur_level = maybe_layer->layerId();
  290. if (!isBatchedAtLevel(tensors, cur_level)) {
  291. return at::_ops::align_tensors::call(tensors);
  292. }
  293. auto results = batch_rule(tensors);
  294. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  295. }
  296. template <typename batch_rule_t, batch_rule_t batch_rule>
  297. void _assert_async_generated_plumbing(const at::Tensor & self) {
  298. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  299. auto maybe_layer = maybeCurrentDynamicLayer();
  300. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  301. int64_t cur_level = maybe_layer->layerId();
  302. if (!isBatchedAtLevel(self, cur_level)) {
  303. return at::_ops::_assert_async::call(self);
  304. }
  305. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  306. batch_rule(self_value, self_bdim);
  307. }
  308. template <typename batch_rule_t, batch_rule_t batch_rule>
  309. void _assert_async_msg_generated_plumbing(const at::Tensor & self, c10::string_view assert_msg) {
  310. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  311. auto maybe_layer = maybeCurrentDynamicLayer();
  312. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  313. int64_t cur_level = maybe_layer->layerId();
  314. if (!isBatchedAtLevel(self, cur_level)) {
  315. return at::_ops::_assert_async_msg::call(self, assert_msg);
  316. }
  317. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  318. batch_rule(self_value, self_bdim, assert_msg);
  319. }
  320. template <typename batch_rule_t, batch_rule_t batch_rule>
  321. at::Tensor _functional_assert_scalar_generated_plumbing(const at::Scalar & self, c10::string_view assert_msg, const at::Tensor & dep_token) {
  322. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  323. auto maybe_layer = maybeCurrentDynamicLayer();
  324. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  325. int64_t cur_level = maybe_layer->layerId();
  326. if (!isBatchedAtLevel(dep_token, cur_level)) {
  327. return at::_ops::_functional_assert_scalar::call(self, assert_msg, dep_token);
  328. }
  329. auto [dep_token_value, dep_token_bdim] = unwrapTensorAtLevel(dep_token, cur_level);
  330. auto results = batch_rule(self, assert_msg, dep_token_value, dep_token_bdim);
  331. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  332. }
  333. template <typename batch_rule_t, batch_rule_t batch_rule>
  334. at::Tensor _functional_assert_async_msg_generated_plumbing(const at::Tensor & self, c10::string_view assert_msg, const at::Tensor & dep_token) {
  335. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  336. auto maybe_layer = maybeCurrentDynamicLayer();
  337. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  338. int64_t cur_level = maybe_layer->layerId();
  339. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(dep_token, cur_level)) {
  340. return at::_ops::_functional_assert_async_msg::call(self, assert_msg, dep_token);
  341. }
  342. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  343. auto [dep_token_value, dep_token_bdim] = unwrapTensorAtLevel(dep_token, cur_level);
  344. auto results = batch_rule(self_value, self_bdim, assert_msg, dep_token_value, dep_token_bdim);
  345. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  346. }
  347. template <typename batch_rule_t, batch_rule_t batch_rule>
  348. void _assert_tensor_metadata_generated_plumbing(const at::Tensor & a, at::OptionalSymIntArrayRef size, at::OptionalSymIntArrayRef stride, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Device> device, ::std::optional<at::Layout> layout) {
  349. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  350. auto maybe_layer = maybeCurrentDynamicLayer();
  351. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  352. int64_t cur_level = maybe_layer->layerId();
  353. if (!isBatchedAtLevel(a, cur_level)) {
  354. return at::_ops::_assert_tensor_metadata::call(a, size, stride, dtype, device, layout);
  355. }
  356. auto [a_value, a_bdim] = unwrapTensorAtLevel(a, cur_level);
  357. batch_rule(a_value, a_bdim, size, stride, dtype, device, layout);
  358. }
  359. template <typename batch_rule_t, batch_rule_t batch_rule>
  360. at::Tensor _functional_sym_constrain_range_generated_plumbing(const at::Scalar & size, ::std::optional<int64_t> min, ::std::optional<int64_t> max, const at::Tensor & dep_token) {
  361. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  362. auto maybe_layer = maybeCurrentDynamicLayer();
  363. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  364. int64_t cur_level = maybe_layer->layerId();
  365. if (!isBatchedAtLevel(dep_token, cur_level)) {
  366. return at::_ops::_functional_sym_constrain_range::call(size, min, max, dep_token);
  367. }
  368. auto [dep_token_value, dep_token_bdim] = unwrapTensorAtLevel(dep_token, cur_level);
  369. auto results = batch_rule(size, min, max, dep_token_value, dep_token_bdim);
  370. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  371. }
  372. template <typename batch_rule_t, batch_rule_t batch_rule>
  373. at::Tensor _functional_sym_constrain_range_for_size_generated_plumbing(const at::Scalar & size, ::std::optional<int64_t> min, ::std::optional<int64_t> max, const at::Tensor & dep_token) {
  374. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  375. auto maybe_layer = maybeCurrentDynamicLayer();
  376. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  377. int64_t cur_level = maybe_layer->layerId();
  378. if (!isBatchedAtLevel(dep_token, cur_level)) {
  379. return at::_ops::_functional_sym_constrain_range_for_size::call(size, min, max, dep_token);
  380. }
  381. auto [dep_token_value, dep_token_bdim] = unwrapTensorAtLevel(dep_token, cur_level);
  382. auto results = batch_rule(size, min, max, dep_token_value, dep_token_bdim);
  383. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  384. }
  385. template <typename batch_rule_t, batch_rule_t batch_rule>
  386. at::Tensor refine_names_generated_plumbing(const at::Tensor & self, at::DimnameList names) {
  387. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  388. auto maybe_layer = maybeCurrentDynamicLayer();
  389. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  390. int64_t cur_level = maybe_layer->layerId();
  391. if (!isBatchedAtLevel(self, cur_level)) {
  392. return at::_ops::refine_names::call(self, names);
  393. }
  394. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  395. auto results = batch_rule(self_value, self_bdim, names);
  396. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  397. }
  398. template <typename batch_rule_t, batch_rule_t batch_rule>
  399. ::std::tuple<at::Tensor,at::Tensor> _cudnn_ctc_loss_generated_plumbing(const at::Tensor & log_probs, const at::Tensor & targets, at::IntArrayRef input_lengths, at::IntArrayRef target_lengths, int64_t blank, bool deterministic, bool zero_infinity) {
  400. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  401. auto maybe_layer = maybeCurrentDynamicLayer();
  402. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  403. int64_t cur_level = maybe_layer->layerId();
  404. if (!isBatchedAtLevel(log_probs, cur_level) && !isBatchedAtLevel(targets, cur_level)) {
  405. return at::_ops::_cudnn_ctc_loss::call(log_probs, targets, input_lengths, target_lengths, blank, deterministic, zero_infinity);
  406. }
  407. auto [log_probs_value, log_probs_bdim] = unwrapTensorAtLevel(log_probs, cur_level);
  408. auto [targets_value, targets_bdim] = unwrapTensorAtLevel(targets, cur_level);
  409. auto results = batch_rule(log_probs_value, log_probs_bdim, targets_value, targets_bdim, input_lengths, target_lengths, blank, deterministic, zero_infinity);
  410. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  411. }
  412. template <typename batch_rule_t, batch_rule_t batch_rule>
  413. ::std::tuple<at::Tensor,at::Tensor> _cudnn_ctc_loss_Tensor_generated_plumbing(const at::Tensor & log_probs, const at::Tensor & targets, const at::Tensor & input_lengths, const at::Tensor & target_lengths, int64_t blank, bool deterministic, bool zero_infinity) {
  414. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  415. auto maybe_layer = maybeCurrentDynamicLayer();
  416. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  417. int64_t cur_level = maybe_layer->layerId();
  418. if (!isBatchedAtLevel(log_probs, cur_level) && !isBatchedAtLevel(targets, cur_level) && !isBatchedAtLevel(input_lengths, cur_level) && !isBatchedAtLevel(target_lengths, cur_level)) {
  419. return at::_ops::_cudnn_ctc_loss_Tensor::call(log_probs, targets, input_lengths, target_lengths, blank, deterministic, zero_infinity);
  420. }
  421. auto [log_probs_value, log_probs_bdim] = unwrapTensorAtLevel(log_probs, cur_level);
  422. auto [targets_value, targets_bdim] = unwrapTensorAtLevel(targets, cur_level);
  423. auto [input_lengths_value, input_lengths_bdim] = unwrapTensorAtLevel(input_lengths, cur_level);
  424. auto [target_lengths_value, target_lengths_bdim] = unwrapTensorAtLevel(target_lengths, cur_level);
  425. auto results = batch_rule(log_probs_value, log_probs_bdim, targets_value, targets_bdim, input_lengths_value, input_lengths_bdim, target_lengths_value, target_lengths_bdim, blank, deterministic, zero_infinity);
  426. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  427. }
  428. template <typename batch_rule_t, batch_rule_t batch_rule>
  429. at::Tensor _cudnn_rnn_flatten_weight_generated_plumbing(at::TensorList weight_arr, int64_t weight_stride0, c10::SymInt input_size, int64_t mode, c10::SymInt hidden_size, c10::SymInt proj_size, int64_t num_layers, bool batch_first, bool bidirectional) {
  430. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  431. auto maybe_layer = maybeCurrentDynamicLayer();
  432. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  433. int64_t cur_level = maybe_layer->layerId();
  434. if (!isBatchedAtLevel(weight_arr, cur_level)) {
  435. return at::_ops::_cudnn_rnn_flatten_weight::call(weight_arr, weight_stride0, input_size, mode, hidden_size, proj_size, num_layers, batch_first, bidirectional);
  436. }
  437. auto results = batch_rule(weight_arr, weight_stride0, input_size, mode, hidden_size, proj_size, num_layers, batch_first, bidirectional);
  438. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  439. }
  440. template <typename batch_rule_t, batch_rule_t batch_rule>
  441. ::std::tuple<at::Tensor,at::Tensor,at::Tensor,at::Tensor,at::Tensor> _cudnn_rnn_generated_plumbing(const at::Tensor & input, at::TensorList weight, int64_t weight_stride0, const ::std::optional<at::Tensor> & weight_buf, const at::Tensor & hx, const ::std::optional<at::Tensor> & cx, int64_t mode, c10::SymInt hidden_size, c10::SymInt proj_size, int64_t num_layers, bool batch_first, double dropout, bool train, bool bidirectional, c10::SymIntArrayRef batch_sizes, const ::std::optional<at::Tensor> & dropout_state) {
  442. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  443. auto maybe_layer = maybeCurrentDynamicLayer();
  444. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  445. int64_t cur_level = maybe_layer->layerId();
  446. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(weight_buf, cur_level) && !isBatchedAtLevel(hx, cur_level) && !isBatchedAtLevel(cx, cur_level) && !isBatchedAtLevel(dropout_state, cur_level)) {
  447. return at::_ops::_cudnn_rnn::call(input, weight, weight_stride0, weight_buf, hx, cx, mode, hidden_size, proj_size, num_layers, batch_first, dropout, train, bidirectional, batch_sizes, dropout_state);
  448. }
  449. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  450. auto [hx_value, hx_bdim] = unwrapTensorAtLevel(hx, cur_level);
  451. std::optional<Tensor> weight_buf_value;
  452. std::optional<int64_t> weight_buf_bdim;
  453. if (weight_buf) {
  454. std::tie(weight_buf_value, weight_buf_bdim) = unwrapTensorAtLevel(weight_buf.value(), cur_level);
  455. }
  456. std::optional<Tensor> cx_value;
  457. std::optional<int64_t> cx_bdim;
  458. if (cx) {
  459. std::tie(cx_value, cx_bdim) = unwrapTensorAtLevel(cx.value(), cur_level);
  460. }
  461. std::optional<Tensor> dropout_state_value;
  462. std::optional<int64_t> dropout_state_bdim;
  463. if (dropout_state) {
  464. std::tie(dropout_state_value, dropout_state_bdim) = unwrapTensorAtLevel(dropout_state.value(), cur_level);
  465. }
  466. auto results = batch_rule(input_value, input_bdim, weight, weight_stride0, weight_buf_value, weight_buf_bdim, hx_value, hx_bdim, cx_value, cx_bdim, mode, hidden_size, proj_size, num_layers, batch_first, dropout, train, bidirectional, batch_sizes, dropout_state_value, dropout_state_bdim);
  467. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level), makeBatched(std::get<6>(results), std::get<7>(results), cur_level), makeBatched(std::get<8>(results), std::get<9>(results), cur_level));
  468. }
  469. template <typename batch_rule_t, batch_rule_t batch_rule>
  470. ::std::tuple<at::Tensor,at::Tensor,at::Tensor,::std::vector<at::Tensor>> _cudnn_rnn_backward_generated_plumbing(const at::Tensor & input, at::TensorList weight, int64_t weight_stride0, const at::Tensor & weight_buf, const at::Tensor & hx, const ::std::optional<at::Tensor> & cx, const at::Tensor & output, const ::std::optional<at::Tensor> & grad_output, const ::std::optional<at::Tensor> & grad_hy, const ::std::optional<at::Tensor> & grad_cy, int64_t mode, c10::SymInt hidden_size, c10::SymInt proj_size, int64_t num_layers, bool batch_first, double dropout, bool train, bool bidirectional, c10::SymIntArrayRef batch_sizes, const ::std::optional<at::Tensor> & dropout_state, const at::Tensor & reserve, ::std::array<bool,4> output_mask) {
  471. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  472. auto maybe_layer = maybeCurrentDynamicLayer();
  473. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  474. int64_t cur_level = maybe_layer->layerId();
  475. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(weight_buf, cur_level) && !isBatchedAtLevel(hx, cur_level) && !isBatchedAtLevel(cx, cur_level) && !isBatchedAtLevel(output, cur_level) && !isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(grad_hy, cur_level) && !isBatchedAtLevel(grad_cy, cur_level) && !isBatchedAtLevel(dropout_state, cur_level) && !isBatchedAtLevel(reserve, cur_level)) {
  476. return at::_ops::_cudnn_rnn_backward::call(input, weight, weight_stride0, weight_buf, hx, cx, output, grad_output, grad_hy, grad_cy, mode, hidden_size, proj_size, num_layers, batch_first, dropout, train, bidirectional, batch_sizes, dropout_state, reserve, output_mask);
  477. }
  478. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  479. auto [weight_buf_value, weight_buf_bdim] = unwrapTensorAtLevel(weight_buf, cur_level);
  480. auto [hx_value, hx_bdim] = unwrapTensorAtLevel(hx, cur_level);
  481. auto [output_value, output_bdim] = unwrapTensorAtLevel(output, cur_level);
  482. auto [reserve_value, reserve_bdim] = unwrapTensorAtLevel(reserve, cur_level);
  483. std::optional<Tensor> cx_value;
  484. std::optional<int64_t> cx_bdim;
  485. if (cx) {
  486. std::tie(cx_value, cx_bdim) = unwrapTensorAtLevel(cx.value(), cur_level);
  487. }
  488. std::optional<Tensor> grad_output_value;
  489. std::optional<int64_t> grad_output_bdim;
  490. if (grad_output) {
  491. std::tie(grad_output_value, grad_output_bdim) = unwrapTensorAtLevel(grad_output.value(), cur_level);
  492. }
  493. std::optional<Tensor> grad_hy_value;
  494. std::optional<int64_t> grad_hy_bdim;
  495. if (grad_hy) {
  496. std::tie(grad_hy_value, grad_hy_bdim) = unwrapTensorAtLevel(grad_hy.value(), cur_level);
  497. }
  498. std::optional<Tensor> grad_cy_value;
  499. std::optional<int64_t> grad_cy_bdim;
  500. if (grad_cy) {
  501. std::tie(grad_cy_value, grad_cy_bdim) = unwrapTensorAtLevel(grad_cy.value(), cur_level);
  502. }
  503. std::optional<Tensor> dropout_state_value;
  504. std::optional<int64_t> dropout_state_bdim;
  505. if (dropout_state) {
  506. std::tie(dropout_state_value, dropout_state_bdim) = unwrapTensorAtLevel(dropout_state.value(), cur_level);
  507. }
  508. auto results = batch_rule(input_value, input_bdim, weight, weight_stride0, weight_buf_value, weight_buf_bdim, hx_value, hx_bdim, cx_value, cx_bdim, output_value, output_bdim, grad_output_value, grad_output_bdim, grad_hy_value, grad_hy_bdim, grad_cy_value, grad_cy_bdim, mode, hidden_size, proj_size, num_layers, batch_first, dropout, train, bidirectional, batch_sizes, dropout_state_value, dropout_state_bdim, reserve_value, reserve_bdim, output_mask);
  509. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level), makeBatchedVector(std::get<6>(results), std::get<7>(results), cur_level));
  510. }
  511. template <typename batch_rule_t, batch_rule_t batch_rule>
  512. ::std::tuple<at::Tensor,at::Tensor> _fused_dropout_generated_plumbing(const at::Tensor & self, double p, ::std::optional<at::Generator> generator) {
  513. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  514. auto maybe_layer = maybeCurrentDynamicLayer();
  515. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  516. int64_t cur_level = maybe_layer->layerId();
  517. if (!isBatchedAtLevel(self, cur_level)) {
  518. return at::_ops::_fused_dropout::call(self, p, generator);
  519. }
  520. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  521. auto results = batch_rule(self_value, self_bdim, p, generator);
  522. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  523. }
  524. template <typename batch_rule_t, batch_rule_t batch_rule>
  525. at::Tensor _masked_scale_generated_plumbing(const at::Tensor & self, const at::Tensor & mask, double scale) {
  526. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  527. auto maybe_layer = maybeCurrentDynamicLayer();
  528. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  529. int64_t cur_level = maybe_layer->layerId();
  530. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(mask, cur_level)) {
  531. return at::_ops::_masked_scale::call(self, mask, scale);
  532. }
  533. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  534. auto [mask_value, mask_bdim] = unwrapTensorAtLevel(mask, cur_level);
  535. auto results = batch_rule(self_value, self_bdim, mask_value, mask_bdim, scale);
  536. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  537. }
  538. template <typename batch_rule_t, batch_rule_t batch_rule>
  539. ::std::tuple<at::Tensor,at::Tensor> native_dropout_generated_plumbing(const at::Tensor & input, double p, ::std::optional<bool> train) {
  540. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  541. auto maybe_layer = maybeCurrentDynamicLayer();
  542. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  543. int64_t cur_level = maybe_layer->layerId();
  544. if (!isBatchedAtLevel(input, cur_level)) {
  545. return at::_ops::native_dropout::call(input, p, train);
  546. }
  547. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  548. auto results = batch_rule(input_value, input_bdim, p, train);
  549. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  550. }
  551. template <typename batch_rule_t, batch_rule_t batch_rule>
  552. at::Tensor native_dropout_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & mask, double scale) {
  553. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  554. auto maybe_layer = maybeCurrentDynamicLayer();
  555. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  556. int64_t cur_level = maybe_layer->layerId();
  557. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(mask, cur_level)) {
  558. return at::_ops::native_dropout_backward::call(grad_output, mask, scale);
  559. }
  560. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  561. auto [mask_value, mask_bdim] = unwrapTensorAtLevel(mask, cur_level);
  562. auto results = batch_rule(grad_output_value, grad_output_bdim, mask_value, mask_bdim, scale);
  563. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  564. }
  565. template <typename batch_rule_t, batch_rule_t batch_rule>
  566. ::std::tuple<at::Tensor,at::Tensor> _sobol_engine_draw_generated_plumbing(const at::Tensor & quasi, int64_t n, const at::Tensor & sobolstate, int64_t dimension, int64_t num_generated, ::std::optional<at::ScalarType> dtype) {
  567. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  568. auto maybe_layer = maybeCurrentDynamicLayer();
  569. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  570. int64_t cur_level = maybe_layer->layerId();
  571. if (!isBatchedAtLevel(quasi, cur_level) && !isBatchedAtLevel(sobolstate, cur_level)) {
  572. return at::_ops::_sobol_engine_draw::call(quasi, n, sobolstate, dimension, num_generated, dtype);
  573. }
  574. auto [quasi_value, quasi_bdim] = unwrapTensorAtLevel(quasi, cur_level);
  575. auto [sobolstate_value, sobolstate_bdim] = unwrapTensorAtLevel(sobolstate, cur_level);
  576. auto results = batch_rule(quasi_value, quasi_bdim, n, sobolstate_value, sobolstate_bdim, dimension, num_generated, dtype);
  577. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  578. }
  579. template <typename batch_rule_t, batch_rule_t batch_rule>
  580. at::Tensor & _sobol_engine_ff__generated_plumbing(at::Tensor & self, int64_t n, const at::Tensor & sobolstate, int64_t dimension, int64_t num_generated) {
  581. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  582. auto maybe_layer = maybeCurrentDynamicLayer();
  583. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  584. int64_t cur_level = maybe_layer->layerId();
  585. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(sobolstate, cur_level)) {
  586. return at::_ops::_sobol_engine_ff_::call(self, n, sobolstate, dimension, num_generated);
  587. }
  588. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  589. auto [sobolstate_value, sobolstate_bdim] = unwrapTensorAtLevel(sobolstate, cur_level);
  590. batch_rule(self_value, self_bdim, n, sobolstate_value, sobolstate_bdim, dimension, num_generated);
  591. return self;
  592. }
  593. template <typename batch_rule_t, batch_rule_t batch_rule>
  594. at::Tensor & _sobol_engine_scramble__generated_plumbing(at::Tensor & self, const at::Tensor & ltm, int64_t dimension) {
  595. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  596. auto maybe_layer = maybeCurrentDynamicLayer();
  597. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  598. int64_t cur_level = maybe_layer->layerId();
  599. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(ltm, cur_level)) {
  600. return at::_ops::_sobol_engine_scramble_::call(self, ltm, dimension);
  601. }
  602. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  603. auto [ltm_value, ltm_bdim] = unwrapTensorAtLevel(ltm, cur_level);
  604. batch_rule(self_value, self_bdim, ltm_value, ltm_bdim, dimension);
  605. return self;
  606. }
  607. template <typename batch_rule_t, batch_rule_t batch_rule>
  608. at::Tensor & _sobol_engine_initialize_state__generated_plumbing(at::Tensor & self, int64_t dimension) {
  609. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  610. auto maybe_layer = maybeCurrentDynamicLayer();
  611. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  612. int64_t cur_level = maybe_layer->layerId();
  613. if (!isBatchedAtLevel(self, cur_level)) {
  614. return at::_ops::_sobol_engine_initialize_state_::call(self, dimension);
  615. }
  616. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  617. batch_rule(self_value, self_bdim, dimension);
  618. return self;
  619. }
  620. template <typename batch_rule_t, batch_rule_t batch_rule>
  621. at::Tensor _reshape_from_tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & shape) {
  622. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  623. auto maybe_layer = maybeCurrentDynamicLayer();
  624. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  625. int64_t cur_level = maybe_layer->layerId();
  626. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(shape, cur_level)) {
  627. return at::_ops::_reshape_from_tensor::call(self, shape);
  628. }
  629. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  630. auto [shape_value, shape_bdim] = unwrapTensorAtLevel(shape, cur_level);
  631. auto results = batch_rule(self_value, self_bdim, shape_value, shape_bdim);
  632. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  633. }
  634. template <typename batch_rule_t, batch_rule_t batch_rule>
  635. at::Tensor _shape_as_tensor_generated_plumbing(const at::Tensor & self) {
  636. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  637. auto maybe_layer = maybeCurrentDynamicLayer();
  638. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  639. int64_t cur_level = maybe_layer->layerId();
  640. if (!isBatchedAtLevel(self, cur_level)) {
  641. return at::_ops::_shape_as_tensor::call(self);
  642. }
  643. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  644. auto results = batch_rule(self_value, self_bdim);
  645. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  646. }
  647. template <typename batch_rule_t, batch_rule_t batch_rule>
  648. at::Tensor dropout_generated_plumbing(const at::Tensor & input, double p, bool train) {
  649. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  650. auto maybe_layer = maybeCurrentDynamicLayer();
  651. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  652. int64_t cur_level = maybe_layer->layerId();
  653. if (!isBatchedAtLevel(input, cur_level)) {
  654. return at::_ops::dropout::call(input, p, train);
  655. }
  656. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  657. auto results = batch_rule(input_value, input_bdim, p, train);
  658. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  659. }
  660. template <typename batch_rule_t, batch_rule_t batch_rule>
  661. at::Tensor & dropout__generated_plumbing(at::Tensor & self, double p, bool train) {
  662. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  663. auto maybe_layer = maybeCurrentDynamicLayer();
  664. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  665. int64_t cur_level = maybe_layer->layerId();
  666. if (!isBatchedAtLevel(self, cur_level)) {
  667. return at::_ops::dropout_::call(self, p, train);
  668. }
  669. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  670. batch_rule(self_value, self_bdim, p, train);
  671. return self;
  672. }
  673. template <typename batch_rule_t, batch_rule_t batch_rule>
  674. at::Tensor feature_dropout_generated_plumbing(const at::Tensor & input, double p, bool train) {
  675. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  676. auto maybe_layer = maybeCurrentDynamicLayer();
  677. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  678. int64_t cur_level = maybe_layer->layerId();
  679. if (!isBatchedAtLevel(input, cur_level)) {
  680. return at::_ops::feature_dropout::call(input, p, train);
  681. }
  682. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  683. auto results = batch_rule(input_value, input_bdim, p, train);
  684. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  685. }
  686. template <typename batch_rule_t, batch_rule_t batch_rule>
  687. at::Tensor & feature_dropout__generated_plumbing(at::Tensor & self, double p, bool train) {
  688. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  689. auto maybe_layer = maybeCurrentDynamicLayer();
  690. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  691. int64_t cur_level = maybe_layer->layerId();
  692. if (!isBatchedAtLevel(self, cur_level)) {
  693. return at::_ops::feature_dropout_::call(self, p, train);
  694. }
  695. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  696. batch_rule(self_value, self_bdim, p, train);
  697. return self;
  698. }
  699. template <typename batch_rule_t, batch_rule_t batch_rule>
  700. at::Tensor alpha_dropout_generated_plumbing(const at::Tensor & input, double p, bool train) {
  701. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  702. auto maybe_layer = maybeCurrentDynamicLayer();
  703. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  704. int64_t cur_level = maybe_layer->layerId();
  705. if (!isBatchedAtLevel(input, cur_level)) {
  706. return at::_ops::alpha_dropout::call(input, p, train);
  707. }
  708. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  709. auto results = batch_rule(input_value, input_bdim, p, train);
  710. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  711. }
  712. template <typename batch_rule_t, batch_rule_t batch_rule>
  713. at::Tensor & alpha_dropout__generated_plumbing(at::Tensor & self, double p, bool train) {
  714. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  715. auto maybe_layer = maybeCurrentDynamicLayer();
  716. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  717. int64_t cur_level = maybe_layer->layerId();
  718. if (!isBatchedAtLevel(self, cur_level)) {
  719. return at::_ops::alpha_dropout_::call(self, p, train);
  720. }
  721. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  722. batch_rule(self_value, self_bdim, p, train);
  723. return self;
  724. }
  725. template <typename batch_rule_t, batch_rule_t batch_rule>
  726. at::Tensor feature_alpha_dropout_generated_plumbing(const at::Tensor & input, double p, bool train) {
  727. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  728. auto maybe_layer = maybeCurrentDynamicLayer();
  729. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  730. int64_t cur_level = maybe_layer->layerId();
  731. if (!isBatchedAtLevel(input, cur_level)) {
  732. return at::_ops::feature_alpha_dropout::call(input, p, train);
  733. }
  734. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  735. auto results = batch_rule(input_value, input_bdim, p, train);
  736. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  737. }
  738. template <typename batch_rule_t, batch_rule_t batch_rule>
  739. at::Tensor & feature_alpha_dropout__generated_plumbing(at::Tensor & self, double p, bool train) {
  740. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  741. auto maybe_layer = maybeCurrentDynamicLayer();
  742. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  743. int64_t cur_level = maybe_layer->layerId();
  744. if (!isBatchedAtLevel(self, cur_level)) {
  745. return at::_ops::feature_alpha_dropout_::call(self, p, train);
  746. }
  747. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  748. batch_rule(self_value, self_bdim, p, train);
  749. return self;
  750. }
  751. template <typename batch_rule_t, batch_rule_t batch_rule>
  752. at::Tensor abs_generated_plumbing(const at::Tensor & self) {
  753. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  754. auto maybe_layer = maybeCurrentDynamicLayer();
  755. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  756. int64_t cur_level = maybe_layer->layerId();
  757. if (!isBatchedAtLevel(self, cur_level)) {
  758. return at::_ops::abs::call(self);
  759. }
  760. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  761. auto results = batch_rule(self_value, self_bdim);
  762. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  763. }
  764. template <typename batch_rule_t, batch_rule_t batch_rule>
  765. at::Tensor & abs__generated_plumbing(at::Tensor & self) {
  766. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  767. auto maybe_layer = maybeCurrentDynamicLayer();
  768. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  769. int64_t cur_level = maybe_layer->layerId();
  770. if (!isBatchedAtLevel(self, cur_level)) {
  771. return at::_ops::abs_::call(self);
  772. }
  773. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  774. batch_rule(self_value, self_bdim);
  775. return self;
  776. }
  777. template <typename batch_rule_t, batch_rule_t batch_rule>
  778. at::Tensor absolute_generated_plumbing(const at::Tensor & self) {
  779. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  780. auto maybe_layer = maybeCurrentDynamicLayer();
  781. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  782. int64_t cur_level = maybe_layer->layerId();
  783. if (!isBatchedAtLevel(self, cur_level)) {
  784. return at::_ops::absolute::call(self);
  785. }
  786. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  787. auto results = batch_rule(self_value, self_bdim);
  788. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  789. }
  790. template <typename batch_rule_t, batch_rule_t batch_rule>
  791. at::Tensor & absolute__generated_plumbing(at::Tensor & self) {
  792. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  793. auto maybe_layer = maybeCurrentDynamicLayer();
  794. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  795. int64_t cur_level = maybe_layer->layerId();
  796. if (!isBatchedAtLevel(self, cur_level)) {
  797. return at::_ops::absolute_::call(self);
  798. }
  799. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  800. batch_rule(self_value, self_bdim);
  801. return self;
  802. }
  803. template <typename batch_rule_t, batch_rule_t batch_rule>
  804. at::Tensor angle_generated_plumbing(const at::Tensor & self) {
  805. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  806. auto maybe_layer = maybeCurrentDynamicLayer();
  807. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  808. int64_t cur_level = maybe_layer->layerId();
  809. if (!isBatchedAtLevel(self, cur_level)) {
  810. return at::_ops::angle::call(self);
  811. }
  812. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  813. auto results = batch_rule(self_value, self_bdim);
  814. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  815. }
  816. template <typename batch_rule_t, batch_rule_t batch_rule>
  817. at::Tensor view_as_real_generated_plumbing(const at::Tensor & self) {
  818. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  819. auto maybe_layer = maybeCurrentDynamicLayer();
  820. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  821. int64_t cur_level = maybe_layer->layerId();
  822. if (!isBatchedAtLevel(self, cur_level)) {
  823. return at::_ops::view_as_real::call(self);
  824. }
  825. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  826. auto results = batch_rule(self_value, self_bdim);
  827. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  828. }
  829. template <typename batch_rule_t, batch_rule_t batch_rule>
  830. at::Tensor view_as_complex_generated_plumbing(const at::Tensor & self) {
  831. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  832. auto maybe_layer = maybeCurrentDynamicLayer();
  833. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  834. int64_t cur_level = maybe_layer->layerId();
  835. if (!isBatchedAtLevel(self, cur_level)) {
  836. return at::_ops::view_as_complex::call(self);
  837. }
  838. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  839. auto results = batch_rule(self_value, self_bdim);
  840. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  841. }
  842. template <typename batch_rule_t, batch_rule_t batch_rule>
  843. at::Tensor sgn_generated_plumbing(const at::Tensor & self) {
  844. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  845. auto maybe_layer = maybeCurrentDynamicLayer();
  846. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  847. int64_t cur_level = maybe_layer->layerId();
  848. if (!isBatchedAtLevel(self, cur_level)) {
  849. return at::_ops::sgn::call(self);
  850. }
  851. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  852. auto results = batch_rule(self_value, self_bdim);
  853. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  854. }
  855. template <typename batch_rule_t, batch_rule_t batch_rule>
  856. at::Tensor & sgn__generated_plumbing(at::Tensor & self) {
  857. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  858. auto maybe_layer = maybeCurrentDynamicLayer();
  859. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  860. int64_t cur_level = maybe_layer->layerId();
  861. if (!isBatchedAtLevel(self, cur_level)) {
  862. return at::_ops::sgn_::call(self);
  863. }
  864. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  865. batch_rule(self_value, self_bdim);
  866. return self;
  867. }
  868. template <typename batch_rule_t, batch_rule_t batch_rule>
  869. at::Tensor chalf_generated_plumbing(const at::Tensor & self, ::std::optional<at::MemoryFormat> memory_format) {
  870. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  871. auto maybe_layer = maybeCurrentDynamicLayer();
  872. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  873. int64_t cur_level = maybe_layer->layerId();
  874. if (!isBatchedAtLevel(self, cur_level)) {
  875. return at::_ops::chalf::call(self, memory_format);
  876. }
  877. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  878. auto results = batch_rule(self_value, self_bdim, memory_format);
  879. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  880. }
  881. template <typename batch_rule_t, batch_rule_t batch_rule>
  882. at::Tensor real_generated_plumbing(const at::Tensor & self) {
  883. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  884. auto maybe_layer = maybeCurrentDynamicLayer();
  885. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  886. int64_t cur_level = maybe_layer->layerId();
  887. if (!isBatchedAtLevel(self, cur_level)) {
  888. return at::_ops::real::call(self);
  889. }
  890. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  891. auto results = batch_rule(self_value, self_bdim);
  892. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  893. }
  894. template <typename batch_rule_t, batch_rule_t batch_rule>
  895. at::Tensor imag_generated_plumbing(const at::Tensor & self) {
  896. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  897. auto maybe_layer = maybeCurrentDynamicLayer();
  898. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  899. int64_t cur_level = maybe_layer->layerId();
  900. if (!isBatchedAtLevel(self, cur_level)) {
  901. return at::_ops::imag::call(self);
  902. }
  903. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  904. auto results = batch_rule(self_value, self_bdim);
  905. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  906. }
  907. template <typename batch_rule_t, batch_rule_t batch_rule>
  908. at::Tensor _conj_generated_plumbing(const at::Tensor & self) {
  909. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  910. auto maybe_layer = maybeCurrentDynamicLayer();
  911. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  912. int64_t cur_level = maybe_layer->layerId();
  913. if (!isBatchedAtLevel(self, cur_level)) {
  914. return at::_ops::_conj::call(self);
  915. }
  916. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  917. auto results = batch_rule(self_value, self_bdim);
  918. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  919. }
  920. template <typename batch_rule_t, batch_rule_t batch_rule>
  921. at::Tensor conj_generated_plumbing(const at::Tensor & self) {
  922. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  923. auto maybe_layer = maybeCurrentDynamicLayer();
  924. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  925. int64_t cur_level = maybe_layer->layerId();
  926. if (!isBatchedAtLevel(self, cur_level)) {
  927. return at::_ops::conj::call(self);
  928. }
  929. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  930. auto results = batch_rule(self_value, self_bdim);
  931. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  932. }
  933. template <typename batch_rule_t, batch_rule_t batch_rule>
  934. at::Tensor _conj_physical_generated_plumbing(const at::Tensor & self) {
  935. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  936. auto maybe_layer = maybeCurrentDynamicLayer();
  937. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  938. int64_t cur_level = maybe_layer->layerId();
  939. if (!isBatchedAtLevel(self, cur_level)) {
  940. return at::_ops::_conj_physical::call(self);
  941. }
  942. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  943. auto results = batch_rule(self_value, self_bdim);
  944. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  945. }
  946. template <typename batch_rule_t, batch_rule_t batch_rule>
  947. at::Tensor conj_physical_generated_plumbing(const at::Tensor & self) {
  948. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  949. auto maybe_layer = maybeCurrentDynamicLayer();
  950. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  951. int64_t cur_level = maybe_layer->layerId();
  952. if (!isBatchedAtLevel(self, cur_level)) {
  953. return at::_ops::conj_physical::call(self);
  954. }
  955. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  956. auto results = batch_rule(self_value, self_bdim);
  957. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  958. }
  959. template <typename batch_rule_t, batch_rule_t batch_rule>
  960. at::Tensor & conj_physical__generated_plumbing(at::Tensor & self) {
  961. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  962. auto maybe_layer = maybeCurrentDynamicLayer();
  963. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  964. int64_t cur_level = maybe_layer->layerId();
  965. if (!isBatchedAtLevel(self, cur_level)) {
  966. return at::_ops::conj_physical_::call(self);
  967. }
  968. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  969. batch_rule(self_value, self_bdim);
  970. return self;
  971. }
  972. template <typename batch_rule_t, batch_rule_t batch_rule>
  973. at::Tensor resolve_conj_generated_plumbing(const at::Tensor & self) {
  974. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  975. auto maybe_layer = maybeCurrentDynamicLayer();
  976. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  977. int64_t cur_level = maybe_layer->layerId();
  978. if (!isBatchedAtLevel(self, cur_level)) {
  979. return at::_ops::resolve_conj::call(self);
  980. }
  981. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  982. auto results = batch_rule(self_value, self_bdim);
  983. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  984. }
  985. template <typename batch_rule_t, batch_rule_t batch_rule>
  986. at::Tensor resolve_neg_generated_plumbing(const at::Tensor & self) {
  987. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  988. auto maybe_layer = maybeCurrentDynamicLayer();
  989. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  990. int64_t cur_level = maybe_layer->layerId();
  991. if (!isBatchedAtLevel(self, cur_level)) {
  992. return at::_ops::resolve_neg::call(self);
  993. }
  994. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  995. auto results = batch_rule(self_value, self_bdim);
  996. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  997. }
  998. template <typename batch_rule_t, batch_rule_t batch_rule>
  999. at::Tensor _neg_view_generated_plumbing(const at::Tensor & self) {
  1000. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1001. auto maybe_layer = maybeCurrentDynamicLayer();
  1002. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1003. int64_t cur_level = maybe_layer->layerId();
  1004. if (!isBatchedAtLevel(self, cur_level)) {
  1005. return at::_ops::_neg_view::call(self);
  1006. }
  1007. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1008. auto results = batch_rule(self_value, self_bdim);
  1009. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1010. }
  1011. template <typename batch_rule_t, batch_rule_t batch_rule>
  1012. at::Tensor acos_generated_plumbing(const at::Tensor & self) {
  1013. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1014. auto maybe_layer = maybeCurrentDynamicLayer();
  1015. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1016. int64_t cur_level = maybe_layer->layerId();
  1017. if (!isBatchedAtLevel(self, cur_level)) {
  1018. return at::_ops::acos::call(self);
  1019. }
  1020. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1021. auto results = batch_rule(self_value, self_bdim);
  1022. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1023. }
  1024. template <typename batch_rule_t, batch_rule_t batch_rule>
  1025. at::Tensor & acos__generated_plumbing(at::Tensor & self) {
  1026. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1027. auto maybe_layer = maybeCurrentDynamicLayer();
  1028. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  1029. int64_t cur_level = maybe_layer->layerId();
  1030. if (!isBatchedAtLevel(self, cur_level)) {
  1031. return at::_ops::acos_::call(self);
  1032. }
  1033. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1034. batch_rule(self_value, self_bdim);
  1035. return self;
  1036. }
  1037. template <typename batch_rule_t, batch_rule_t batch_rule>
  1038. at::Tensor arccos_generated_plumbing(const at::Tensor & self) {
  1039. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1040. auto maybe_layer = maybeCurrentDynamicLayer();
  1041. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1042. int64_t cur_level = maybe_layer->layerId();
  1043. if (!isBatchedAtLevel(self, cur_level)) {
  1044. return at::_ops::arccos::call(self);
  1045. }
  1046. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1047. auto results = batch_rule(self_value, self_bdim);
  1048. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1049. }
  1050. template <typename batch_rule_t, batch_rule_t batch_rule>
  1051. at::Tensor & arccos__generated_plumbing(at::Tensor & self) {
  1052. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1053. auto maybe_layer = maybeCurrentDynamicLayer();
  1054. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  1055. int64_t cur_level = maybe_layer->layerId();
  1056. if (!isBatchedAtLevel(self, cur_level)) {
  1057. return at::_ops::arccos_::call(self);
  1058. }
  1059. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1060. batch_rule(self_value, self_bdim);
  1061. return self;
  1062. }
  1063. template <typename batch_rule_t, batch_rule_t batch_rule>
  1064. at::Tensor avg_pool1d_generated_plumbing(const at::Tensor & self, at::IntArrayRef kernel_size, at::IntArrayRef stride, at::IntArrayRef padding, bool ceil_mode, bool count_include_pad) {
  1065. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1066. auto maybe_layer = maybeCurrentDynamicLayer();
  1067. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1068. int64_t cur_level = maybe_layer->layerId();
  1069. if (!isBatchedAtLevel(self, cur_level)) {
  1070. return at::_ops::avg_pool1d::call(self, kernel_size, stride, padding, ceil_mode, count_include_pad);
  1071. }
  1072. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1073. auto results = batch_rule(self_value, self_bdim, kernel_size, stride, padding, ceil_mode, count_include_pad);
  1074. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1075. }
  1076. template <typename batch_rule_t, batch_rule_t batch_rule>
  1077. at::Tensor adaptive_avg_pool1d_generated_plumbing(const at::Tensor & self, at::IntArrayRef output_size) {
  1078. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1079. auto maybe_layer = maybeCurrentDynamicLayer();
  1080. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1081. int64_t cur_level = maybe_layer->layerId();
  1082. if (!isBatchedAtLevel(self, cur_level)) {
  1083. return at::_ops::adaptive_avg_pool1d::call(self, output_size);
  1084. }
  1085. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1086. auto results = batch_rule(self_value, self_bdim, output_size);
  1087. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1088. }
  1089. template <typename batch_rule_t, batch_rule_t batch_rule>
  1090. ::std::tuple<at::Tensor,at::Tensor> adaptive_max_pool1d_generated_plumbing(const at::Tensor & self, at::IntArrayRef output_size) {
  1091. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1092. auto maybe_layer = maybeCurrentDynamicLayer();
  1093. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1094. int64_t cur_level = maybe_layer->layerId();
  1095. if (!isBatchedAtLevel(self, cur_level)) {
  1096. return at::_ops::adaptive_max_pool1d::call(self, output_size);
  1097. }
  1098. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1099. auto results = batch_rule(self_value, self_bdim, output_size);
  1100. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  1101. }
  1102. template <typename batch_rule_t, batch_rule_t batch_rule>
  1103. at::Tensor add_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & other, const at::Scalar & alpha) {
  1104. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1105. auto maybe_layer = maybeCurrentDynamicLayer();
  1106. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1107. int64_t cur_level = maybe_layer->layerId();
  1108. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  1109. return at::_ops::add_Tensor::call(self, other, alpha);
  1110. }
  1111. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1112. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  1113. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim, alpha);
  1114. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1115. }
  1116. template <typename batch_rule_t, batch_rule_t batch_rule>
  1117. at::Tensor & add__Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & other, const at::Scalar & alpha) {
  1118. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1119. auto maybe_layer = maybeCurrentDynamicLayer();
  1120. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  1121. int64_t cur_level = maybe_layer->layerId();
  1122. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  1123. return at::_ops::add__Tensor::call(self, other, alpha);
  1124. }
  1125. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1126. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  1127. batch_rule(self_value, self_bdim, other_value, other_bdim, alpha);
  1128. return self;
  1129. }
  1130. template <typename batch_rule_t, batch_rule_t batch_rule>
  1131. at::Tensor _add_relu_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & other, const at::Scalar & alpha) {
  1132. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1133. auto maybe_layer = maybeCurrentDynamicLayer();
  1134. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1135. int64_t cur_level = maybe_layer->layerId();
  1136. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  1137. return at::_ops::_add_relu_Tensor::call(self, other, alpha);
  1138. }
  1139. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1140. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  1141. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim, alpha);
  1142. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1143. }
  1144. template <typename batch_rule_t, batch_rule_t batch_rule>
  1145. at::Tensor & _add_relu__Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & other, const at::Scalar & alpha) {
  1146. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1147. auto maybe_layer = maybeCurrentDynamicLayer();
  1148. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  1149. int64_t cur_level = maybe_layer->layerId();
  1150. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  1151. return at::_ops::_add_relu__Tensor::call(self, other, alpha);
  1152. }
  1153. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1154. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  1155. batch_rule(self_value, self_bdim, other_value, other_bdim, alpha);
  1156. return self;
  1157. }
  1158. template <typename batch_rule_t, batch_rule_t batch_rule>
  1159. at::Tensor _add_relu_Scalar_generated_plumbing(const at::Tensor & self, const at::Scalar & other, const at::Scalar & alpha) {
  1160. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1161. auto maybe_layer = maybeCurrentDynamicLayer();
  1162. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1163. int64_t cur_level = maybe_layer->layerId();
  1164. if (!isBatchedAtLevel(self, cur_level)) {
  1165. return at::_ops::_add_relu_Scalar::call(self, other, alpha);
  1166. }
  1167. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1168. auto results = batch_rule(self_value, self_bdim, other, alpha);
  1169. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1170. }
  1171. template <typename batch_rule_t, batch_rule_t batch_rule>
  1172. at::Tensor & _add_relu__Scalar_generated_plumbing(at::Tensor & self, const at::Scalar & other, const at::Scalar & alpha) {
  1173. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1174. auto maybe_layer = maybeCurrentDynamicLayer();
  1175. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  1176. int64_t cur_level = maybe_layer->layerId();
  1177. if (!isBatchedAtLevel(self, cur_level)) {
  1178. return at::_ops::_add_relu__Scalar::call(self, other, alpha);
  1179. }
  1180. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1181. batch_rule(self_value, self_bdim, other, alpha);
  1182. return self;
  1183. }
  1184. template <typename batch_rule_t, batch_rule_t batch_rule>
  1185. at::Tensor add_Scalar_generated_plumbing(const at::Tensor & self, const at::Scalar & other, const at::Scalar & alpha) {
  1186. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1187. auto maybe_layer = maybeCurrentDynamicLayer();
  1188. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1189. int64_t cur_level = maybe_layer->layerId();
  1190. if (!isBatchedAtLevel(self, cur_level)) {
  1191. return at::_ops::add_Scalar::call(self, other, alpha);
  1192. }
  1193. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1194. auto results = batch_rule(self_value, self_bdim, other, alpha);
  1195. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1196. }
  1197. template <typename batch_rule_t, batch_rule_t batch_rule>
  1198. at::Tensor & add__Scalar_generated_plumbing(at::Tensor & self, const at::Scalar & other, const at::Scalar & alpha) {
  1199. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1200. auto maybe_layer = maybeCurrentDynamicLayer();
  1201. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  1202. int64_t cur_level = maybe_layer->layerId();
  1203. if (!isBatchedAtLevel(self, cur_level)) {
  1204. return at::_ops::add__Scalar::call(self, other, alpha);
  1205. }
  1206. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1207. batch_rule(self_value, self_bdim, other, alpha);
  1208. return self;
  1209. }
  1210. template <typename batch_rule_t, batch_rule_t batch_rule>
  1211. at::Tensor addmv_generated_plumbing(const at::Tensor & self, const at::Tensor & mat, const at::Tensor & vec, const at::Scalar & beta, const at::Scalar & alpha) {
  1212. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1213. auto maybe_layer = maybeCurrentDynamicLayer();
  1214. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1215. int64_t cur_level = maybe_layer->layerId();
  1216. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(mat, cur_level) && !isBatchedAtLevel(vec, cur_level)) {
  1217. return at::_ops::addmv::call(self, mat, vec, beta, alpha);
  1218. }
  1219. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1220. auto [mat_value, mat_bdim] = unwrapTensorAtLevel(mat, cur_level);
  1221. auto [vec_value, vec_bdim] = unwrapTensorAtLevel(vec, cur_level);
  1222. auto results = batch_rule(self_value, self_bdim, mat_value, mat_bdim, vec_value, vec_bdim, beta, alpha);
  1223. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1224. }
  1225. template <typename batch_rule_t, batch_rule_t batch_rule>
  1226. at::Tensor & addmv__generated_plumbing(at::Tensor & self, const at::Tensor & mat, const at::Tensor & vec, const at::Scalar & beta, const at::Scalar & alpha) {
  1227. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1228. auto maybe_layer = maybeCurrentDynamicLayer();
  1229. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  1230. int64_t cur_level = maybe_layer->layerId();
  1231. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(mat, cur_level) && !isBatchedAtLevel(vec, cur_level)) {
  1232. return at::_ops::addmv_::call(self, mat, vec, beta, alpha);
  1233. }
  1234. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1235. auto [mat_value, mat_bdim] = unwrapTensorAtLevel(mat, cur_level);
  1236. auto [vec_value, vec_bdim] = unwrapTensorAtLevel(vec, cur_level);
  1237. batch_rule(self_value, self_bdim, mat_value, mat_bdim, vec_value, vec_bdim, beta, alpha);
  1238. return self;
  1239. }
  1240. template <typename batch_rule_t, batch_rule_t batch_rule>
  1241. at::Tensor addr_generated_plumbing(const at::Tensor & self, const at::Tensor & vec1, const at::Tensor & vec2, const at::Scalar & beta, const at::Scalar & alpha) {
  1242. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1243. auto maybe_layer = maybeCurrentDynamicLayer();
  1244. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1245. int64_t cur_level = maybe_layer->layerId();
  1246. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(vec1, cur_level) && !isBatchedAtLevel(vec2, cur_level)) {
  1247. return at::_ops::addr::call(self, vec1, vec2, beta, alpha);
  1248. }
  1249. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1250. auto [vec1_value, vec1_bdim] = unwrapTensorAtLevel(vec1, cur_level);
  1251. auto [vec2_value, vec2_bdim] = unwrapTensorAtLevel(vec2, cur_level);
  1252. auto results = batch_rule(self_value, self_bdim, vec1_value, vec1_bdim, vec2_value, vec2_bdim, beta, alpha);
  1253. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1254. }
  1255. template <typename batch_rule_t, batch_rule_t batch_rule>
  1256. at::Tensor & addr__generated_plumbing(at::Tensor & self, const at::Tensor & vec1, const at::Tensor & vec2, const at::Scalar & beta, const at::Scalar & alpha) {
  1257. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1258. auto maybe_layer = maybeCurrentDynamicLayer();
  1259. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  1260. int64_t cur_level = maybe_layer->layerId();
  1261. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(vec1, cur_level) && !isBatchedAtLevel(vec2, cur_level)) {
  1262. return at::_ops::addr_::call(self, vec1, vec2, beta, alpha);
  1263. }
  1264. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1265. auto [vec1_value, vec1_bdim] = unwrapTensorAtLevel(vec1, cur_level);
  1266. auto [vec2_value, vec2_bdim] = unwrapTensorAtLevel(vec2, cur_level);
  1267. batch_rule(self_value, self_bdim, vec1_value, vec1_bdim, vec2_value, vec2_bdim, beta, alpha);
  1268. return self;
  1269. }
  1270. template <typename batch_rule_t, batch_rule_t batch_rule>
  1271. at::Tensor affine_grid_generator_generated_plumbing(const at::Tensor & theta, c10::SymIntArrayRef size, bool align_corners) {
  1272. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1273. auto maybe_layer = maybeCurrentDynamicLayer();
  1274. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1275. int64_t cur_level = maybe_layer->layerId();
  1276. if (!isBatchedAtLevel(theta, cur_level)) {
  1277. return at::_ops::affine_grid_generator::call(theta, size, align_corners);
  1278. }
  1279. auto [theta_value, theta_bdim] = unwrapTensorAtLevel(theta, cur_level);
  1280. auto results = batch_rule(theta_value, theta_bdim, size, align_corners);
  1281. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1282. }
  1283. template <typename batch_rule_t, batch_rule_t batch_rule>
  1284. at::Tensor affine_grid_generator_backward_generated_plumbing(const at::Tensor & grad, c10::SymIntArrayRef size, bool align_corners) {
  1285. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1286. auto maybe_layer = maybeCurrentDynamicLayer();
  1287. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1288. int64_t cur_level = maybe_layer->layerId();
  1289. if (!isBatchedAtLevel(grad, cur_level)) {
  1290. return at::_ops::affine_grid_generator_backward::call(grad, size, align_corners);
  1291. }
  1292. auto [grad_value, grad_bdim] = unwrapTensorAtLevel(grad, cur_level);
  1293. auto results = batch_rule(grad_value, grad_bdim, size, align_corners);
  1294. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1295. }
  1296. template <typename batch_rule_t, batch_rule_t batch_rule>
  1297. at::Tensor _is_all_true_generated_plumbing(const at::Tensor & self) {
  1298. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1299. auto maybe_layer = maybeCurrentDynamicLayer();
  1300. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1301. int64_t cur_level = maybe_layer->layerId();
  1302. if (!isBatchedAtLevel(self, cur_level)) {
  1303. return at::_ops::_is_all_true::call(self);
  1304. }
  1305. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1306. auto results = batch_rule(self_value, self_bdim);
  1307. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1308. }
  1309. template <typename batch_rule_t, batch_rule_t batch_rule>
  1310. at::Tensor _is_any_true_generated_plumbing(const at::Tensor & self) {
  1311. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1312. auto maybe_layer = maybeCurrentDynamicLayer();
  1313. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1314. int64_t cur_level = maybe_layer->layerId();
  1315. if (!isBatchedAtLevel(self, cur_level)) {
  1316. return at::_ops::_is_any_true::call(self);
  1317. }
  1318. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1319. auto results = batch_rule(self_value, self_bdim);
  1320. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1321. }
  1322. template <typename batch_rule_t, batch_rule_t batch_rule>
  1323. at::Tensor _test_check_tensor_generated_plumbing(const at::Tensor & self) {
  1324. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1325. auto maybe_layer = maybeCurrentDynamicLayer();
  1326. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1327. int64_t cur_level = maybe_layer->layerId();
  1328. if (!isBatchedAtLevel(self, cur_level)) {
  1329. return at::_ops::_test_check_tensor::call(self);
  1330. }
  1331. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1332. auto results = batch_rule(self_value, self_bdim);
  1333. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1334. }
  1335. template <typename batch_rule_t, batch_rule_t batch_rule>
  1336. at::Tensor _test_functorch_fallback_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  1337. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1338. auto maybe_layer = maybeCurrentDynamicLayer();
  1339. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1340. int64_t cur_level = maybe_layer->layerId();
  1341. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  1342. return at::_ops::_test_functorch_fallback::call(self, other);
  1343. }
  1344. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1345. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  1346. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  1347. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1348. }
  1349. template <typename batch_rule_t, batch_rule_t batch_rule>
  1350. at::Tensor all_dim_generated_plumbing(const at::Tensor & self, int64_t dim, bool keepdim) {
  1351. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1352. auto maybe_layer = maybeCurrentDynamicLayer();
  1353. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1354. int64_t cur_level = maybe_layer->layerId();
  1355. if (!isBatchedAtLevel(self, cur_level)) {
  1356. return at::_ops::all_dim::call(self, dim, keepdim);
  1357. }
  1358. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1359. auto results = batch_rule(self_value, self_bdim, dim, keepdim);
  1360. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1361. }
  1362. template <typename batch_rule_t, batch_rule_t batch_rule>
  1363. at::Tensor all_dims_generated_plumbing(const at::Tensor & self, at::OptionalIntArrayRef dim, bool keepdim) {
  1364. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1365. auto maybe_layer = maybeCurrentDynamicLayer();
  1366. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1367. int64_t cur_level = maybe_layer->layerId();
  1368. if (!isBatchedAtLevel(self, cur_level)) {
  1369. return at::_ops::all_dims::call(self, dim, keepdim);
  1370. }
  1371. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1372. auto results = batch_rule(self_value, self_bdim, dim, keepdim);
  1373. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1374. }
  1375. template <typename batch_rule_t, batch_rule_t batch_rule>
  1376. at::Tensor all_dimname_generated_plumbing(const at::Tensor & self, at::Dimname dim, bool keepdim) {
  1377. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1378. auto maybe_layer = maybeCurrentDynamicLayer();
  1379. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1380. int64_t cur_level = maybe_layer->layerId();
  1381. if (!isBatchedAtLevel(self, cur_level)) {
  1382. return at::_ops::all_dimname::call(self, dim, keepdim);
  1383. }
  1384. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1385. auto results = batch_rule(self_value, self_bdim, dim, keepdim);
  1386. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1387. }
  1388. template <typename batch_rule_t, batch_rule_t batch_rule>
  1389. at::Tensor any_dim_generated_plumbing(const at::Tensor & self, int64_t dim, bool keepdim) {
  1390. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1391. auto maybe_layer = maybeCurrentDynamicLayer();
  1392. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1393. int64_t cur_level = maybe_layer->layerId();
  1394. if (!isBatchedAtLevel(self, cur_level)) {
  1395. return at::_ops::any_dim::call(self, dim, keepdim);
  1396. }
  1397. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1398. auto results = batch_rule(self_value, self_bdim, dim, keepdim);
  1399. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1400. }
  1401. template <typename batch_rule_t, batch_rule_t batch_rule>
  1402. at::Tensor any_dims_generated_plumbing(const at::Tensor & self, at::OptionalIntArrayRef dim, bool keepdim) {
  1403. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1404. auto maybe_layer = maybeCurrentDynamicLayer();
  1405. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1406. int64_t cur_level = maybe_layer->layerId();
  1407. if (!isBatchedAtLevel(self, cur_level)) {
  1408. return at::_ops::any_dims::call(self, dim, keepdim);
  1409. }
  1410. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1411. auto results = batch_rule(self_value, self_bdim, dim, keepdim);
  1412. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1413. }
  1414. template <typename batch_rule_t, batch_rule_t batch_rule>
  1415. at::Tensor any_dimname_generated_plumbing(const at::Tensor & self, at::Dimname dim, bool keepdim) {
  1416. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1417. auto maybe_layer = maybeCurrentDynamicLayer();
  1418. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1419. int64_t cur_level = maybe_layer->layerId();
  1420. if (!isBatchedAtLevel(self, cur_level)) {
  1421. return at::_ops::any_dimname::call(self, dim, keepdim);
  1422. }
  1423. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1424. auto results = batch_rule(self_value, self_bdim, dim, keepdim);
  1425. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1426. }
  1427. template <typename batch_rule_t, batch_rule_t batch_rule>
  1428. at::Tensor _dim_arange_generated_plumbing(const at::Tensor & like, int64_t dim) {
  1429. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1430. auto maybe_layer = maybeCurrentDynamicLayer();
  1431. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1432. int64_t cur_level = maybe_layer->layerId();
  1433. if (!isBatchedAtLevel(like, cur_level)) {
  1434. return at::_ops::_dim_arange::call(like, dim);
  1435. }
  1436. auto [like_value, like_bdim] = unwrapTensorAtLevel(like, cur_level);
  1437. auto results = batch_rule(like_value, like_bdim, dim);
  1438. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1439. }
  1440. template <typename batch_rule_t, batch_rule_t batch_rule>
  1441. at::Tensor argmax_generated_plumbing(const at::Tensor & self, ::std::optional<int64_t> dim, bool keepdim) {
  1442. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1443. auto maybe_layer = maybeCurrentDynamicLayer();
  1444. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1445. int64_t cur_level = maybe_layer->layerId();
  1446. if (!isBatchedAtLevel(self, cur_level)) {
  1447. return at::_ops::argmax::call(self, dim, keepdim);
  1448. }
  1449. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1450. auto results = batch_rule(self_value, self_bdim, dim, keepdim);
  1451. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1452. }
  1453. template <typename batch_rule_t, batch_rule_t batch_rule>
  1454. at::Tensor argmin_generated_plumbing(const at::Tensor & self, ::std::optional<int64_t> dim, bool keepdim) {
  1455. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1456. auto maybe_layer = maybeCurrentDynamicLayer();
  1457. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1458. int64_t cur_level = maybe_layer->layerId();
  1459. if (!isBatchedAtLevel(self, cur_level)) {
  1460. return at::_ops::argmin::call(self, dim, keepdim);
  1461. }
  1462. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1463. auto results = batch_rule(self_value, self_bdim, dim, keepdim);
  1464. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1465. }
  1466. template <typename batch_rule_t, batch_rule_t batch_rule>
  1467. at::Tensor acosh_generated_plumbing(const at::Tensor & self) {
  1468. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1469. auto maybe_layer = maybeCurrentDynamicLayer();
  1470. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1471. int64_t cur_level = maybe_layer->layerId();
  1472. if (!isBatchedAtLevel(self, cur_level)) {
  1473. return at::_ops::acosh::call(self);
  1474. }
  1475. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1476. auto results = batch_rule(self_value, self_bdim);
  1477. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1478. }
  1479. template <typename batch_rule_t, batch_rule_t batch_rule>
  1480. at::Tensor & acosh__generated_plumbing(at::Tensor & self) {
  1481. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1482. auto maybe_layer = maybeCurrentDynamicLayer();
  1483. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  1484. int64_t cur_level = maybe_layer->layerId();
  1485. if (!isBatchedAtLevel(self, cur_level)) {
  1486. return at::_ops::acosh_::call(self);
  1487. }
  1488. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1489. batch_rule(self_value, self_bdim);
  1490. return self;
  1491. }
  1492. template <typename batch_rule_t, batch_rule_t batch_rule>
  1493. at::Tensor arccosh_generated_plumbing(const at::Tensor & self) {
  1494. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1495. auto maybe_layer = maybeCurrentDynamicLayer();
  1496. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1497. int64_t cur_level = maybe_layer->layerId();
  1498. if (!isBatchedAtLevel(self, cur_level)) {
  1499. return at::_ops::arccosh::call(self);
  1500. }
  1501. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1502. auto results = batch_rule(self_value, self_bdim);
  1503. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1504. }
  1505. template <typename batch_rule_t, batch_rule_t batch_rule>
  1506. at::Tensor & arccosh__generated_plumbing(at::Tensor & self) {
  1507. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1508. auto maybe_layer = maybeCurrentDynamicLayer();
  1509. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  1510. int64_t cur_level = maybe_layer->layerId();
  1511. if (!isBatchedAtLevel(self, cur_level)) {
  1512. return at::_ops::arccosh_::call(self);
  1513. }
  1514. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1515. batch_rule(self_value, self_bdim);
  1516. return self;
  1517. }
  1518. template <typename batch_rule_t, batch_rule_t batch_rule>
  1519. at::Tensor asinh_generated_plumbing(const at::Tensor & self) {
  1520. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1521. auto maybe_layer = maybeCurrentDynamicLayer();
  1522. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1523. int64_t cur_level = maybe_layer->layerId();
  1524. if (!isBatchedAtLevel(self, cur_level)) {
  1525. return at::_ops::asinh::call(self);
  1526. }
  1527. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1528. auto results = batch_rule(self_value, self_bdim);
  1529. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1530. }
  1531. template <typename batch_rule_t, batch_rule_t batch_rule>
  1532. at::Tensor & asinh__generated_plumbing(at::Tensor & self) {
  1533. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1534. auto maybe_layer = maybeCurrentDynamicLayer();
  1535. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  1536. int64_t cur_level = maybe_layer->layerId();
  1537. if (!isBatchedAtLevel(self, cur_level)) {
  1538. return at::_ops::asinh_::call(self);
  1539. }
  1540. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1541. batch_rule(self_value, self_bdim);
  1542. return self;
  1543. }
  1544. template <typename batch_rule_t, batch_rule_t batch_rule>
  1545. at::Tensor arcsinh_generated_plumbing(const at::Tensor & self) {
  1546. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1547. auto maybe_layer = maybeCurrentDynamicLayer();
  1548. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1549. int64_t cur_level = maybe_layer->layerId();
  1550. if (!isBatchedAtLevel(self, cur_level)) {
  1551. return at::_ops::arcsinh::call(self);
  1552. }
  1553. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1554. auto results = batch_rule(self_value, self_bdim);
  1555. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1556. }
  1557. template <typename batch_rule_t, batch_rule_t batch_rule>
  1558. at::Tensor & arcsinh__generated_plumbing(at::Tensor & self) {
  1559. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1560. auto maybe_layer = maybeCurrentDynamicLayer();
  1561. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  1562. int64_t cur_level = maybe_layer->layerId();
  1563. if (!isBatchedAtLevel(self, cur_level)) {
  1564. return at::_ops::arcsinh_::call(self);
  1565. }
  1566. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1567. batch_rule(self_value, self_bdim);
  1568. return self;
  1569. }
  1570. template <typename batch_rule_t, batch_rule_t batch_rule>
  1571. at::Tensor atanh_generated_plumbing(const at::Tensor & self) {
  1572. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1573. auto maybe_layer = maybeCurrentDynamicLayer();
  1574. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1575. int64_t cur_level = maybe_layer->layerId();
  1576. if (!isBatchedAtLevel(self, cur_level)) {
  1577. return at::_ops::atanh::call(self);
  1578. }
  1579. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1580. auto results = batch_rule(self_value, self_bdim);
  1581. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1582. }
  1583. template <typename batch_rule_t, batch_rule_t batch_rule>
  1584. at::Tensor & atanh__generated_plumbing(at::Tensor & self) {
  1585. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1586. auto maybe_layer = maybeCurrentDynamicLayer();
  1587. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  1588. int64_t cur_level = maybe_layer->layerId();
  1589. if (!isBatchedAtLevel(self, cur_level)) {
  1590. return at::_ops::atanh_::call(self);
  1591. }
  1592. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1593. batch_rule(self_value, self_bdim);
  1594. return self;
  1595. }
  1596. template <typename batch_rule_t, batch_rule_t batch_rule>
  1597. at::Tensor arctanh_generated_plumbing(const at::Tensor & self) {
  1598. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1599. auto maybe_layer = maybeCurrentDynamicLayer();
  1600. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1601. int64_t cur_level = maybe_layer->layerId();
  1602. if (!isBatchedAtLevel(self, cur_level)) {
  1603. return at::_ops::arctanh::call(self);
  1604. }
  1605. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1606. auto results = batch_rule(self_value, self_bdim);
  1607. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1608. }
  1609. template <typename batch_rule_t, batch_rule_t batch_rule>
  1610. at::Tensor & arctanh__generated_plumbing(at::Tensor & self) {
  1611. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1612. auto maybe_layer = maybeCurrentDynamicLayer();
  1613. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  1614. int64_t cur_level = maybe_layer->layerId();
  1615. if (!isBatchedAtLevel(self, cur_level)) {
  1616. return at::_ops::arctanh_::call(self);
  1617. }
  1618. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1619. batch_rule(self_value, self_bdim);
  1620. return self;
  1621. }
  1622. template <typename batch_rule_t, batch_rule_t batch_rule>
  1623. at::Tensor as_strided_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef size, c10::SymIntArrayRef stride, ::std::optional<c10::SymInt> storage_offset) {
  1624. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1625. auto maybe_layer = maybeCurrentDynamicLayer();
  1626. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1627. int64_t cur_level = maybe_layer->layerId();
  1628. if (!isBatchedAtLevel(self, cur_level)) {
  1629. return at::_ops::as_strided::call(self, size, stride, storage_offset);
  1630. }
  1631. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1632. auto results = batch_rule(self_value, self_bdim, size, stride, storage_offset);
  1633. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1634. }
  1635. template <typename batch_rule_t, batch_rule_t batch_rule>
  1636. at::Tensor asin_generated_plumbing(const at::Tensor & self) {
  1637. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1638. auto maybe_layer = maybeCurrentDynamicLayer();
  1639. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1640. int64_t cur_level = maybe_layer->layerId();
  1641. if (!isBatchedAtLevel(self, cur_level)) {
  1642. return at::_ops::asin::call(self);
  1643. }
  1644. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1645. auto results = batch_rule(self_value, self_bdim);
  1646. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1647. }
  1648. template <typename batch_rule_t, batch_rule_t batch_rule>
  1649. at::Tensor & asin__generated_plumbing(at::Tensor & self) {
  1650. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1651. auto maybe_layer = maybeCurrentDynamicLayer();
  1652. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  1653. int64_t cur_level = maybe_layer->layerId();
  1654. if (!isBatchedAtLevel(self, cur_level)) {
  1655. return at::_ops::asin_::call(self);
  1656. }
  1657. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1658. batch_rule(self_value, self_bdim);
  1659. return self;
  1660. }
  1661. template <typename batch_rule_t, batch_rule_t batch_rule>
  1662. at::Tensor arcsin_generated_plumbing(const at::Tensor & self) {
  1663. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1664. auto maybe_layer = maybeCurrentDynamicLayer();
  1665. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1666. int64_t cur_level = maybe_layer->layerId();
  1667. if (!isBatchedAtLevel(self, cur_level)) {
  1668. return at::_ops::arcsin::call(self);
  1669. }
  1670. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1671. auto results = batch_rule(self_value, self_bdim);
  1672. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1673. }
  1674. template <typename batch_rule_t, batch_rule_t batch_rule>
  1675. at::Tensor & arcsin__generated_plumbing(at::Tensor & self) {
  1676. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1677. auto maybe_layer = maybeCurrentDynamicLayer();
  1678. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  1679. int64_t cur_level = maybe_layer->layerId();
  1680. if (!isBatchedAtLevel(self, cur_level)) {
  1681. return at::_ops::arcsin_::call(self);
  1682. }
  1683. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1684. batch_rule(self_value, self_bdim);
  1685. return self;
  1686. }
  1687. template <typename batch_rule_t, batch_rule_t batch_rule>
  1688. at::Tensor atan_generated_plumbing(const at::Tensor & self) {
  1689. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1690. auto maybe_layer = maybeCurrentDynamicLayer();
  1691. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1692. int64_t cur_level = maybe_layer->layerId();
  1693. if (!isBatchedAtLevel(self, cur_level)) {
  1694. return at::_ops::atan::call(self);
  1695. }
  1696. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1697. auto results = batch_rule(self_value, self_bdim);
  1698. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1699. }
  1700. template <typename batch_rule_t, batch_rule_t batch_rule>
  1701. at::Tensor & atan__generated_plumbing(at::Tensor & self) {
  1702. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1703. auto maybe_layer = maybeCurrentDynamicLayer();
  1704. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  1705. int64_t cur_level = maybe_layer->layerId();
  1706. if (!isBatchedAtLevel(self, cur_level)) {
  1707. return at::_ops::atan_::call(self);
  1708. }
  1709. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1710. batch_rule(self_value, self_bdim);
  1711. return self;
  1712. }
  1713. template <typename batch_rule_t, batch_rule_t batch_rule>
  1714. at::Tensor arctan_generated_plumbing(const at::Tensor & self) {
  1715. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1716. auto maybe_layer = maybeCurrentDynamicLayer();
  1717. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1718. int64_t cur_level = maybe_layer->layerId();
  1719. if (!isBatchedAtLevel(self, cur_level)) {
  1720. return at::_ops::arctan::call(self);
  1721. }
  1722. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1723. auto results = batch_rule(self_value, self_bdim);
  1724. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1725. }
  1726. template <typename batch_rule_t, batch_rule_t batch_rule>
  1727. at::Tensor & arctan__generated_plumbing(at::Tensor & self) {
  1728. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1729. auto maybe_layer = maybeCurrentDynamicLayer();
  1730. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  1731. int64_t cur_level = maybe_layer->layerId();
  1732. if (!isBatchedAtLevel(self, cur_level)) {
  1733. return at::_ops::arctan_::call(self);
  1734. }
  1735. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1736. batch_rule(self_value, self_bdim);
  1737. return self;
  1738. }
  1739. template <typename batch_rule_t, batch_rule_t batch_rule>
  1740. at::Tensor atleast_1d_generated_plumbing(const at::Tensor & self) {
  1741. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1742. auto maybe_layer = maybeCurrentDynamicLayer();
  1743. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1744. int64_t cur_level = maybe_layer->layerId();
  1745. if (!isBatchedAtLevel(self, cur_level)) {
  1746. return at::_ops::atleast_1d::call(self);
  1747. }
  1748. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1749. auto results = batch_rule(self_value, self_bdim);
  1750. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1751. }
  1752. template <typename batch_rule_t, batch_rule_t batch_rule>
  1753. ::std::vector<at::Tensor> atleast_1d_Sequence_generated_plumbing(at::TensorList tensors) {
  1754. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1755. auto maybe_layer = maybeCurrentDynamicLayer();
  1756. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1757. int64_t cur_level = maybe_layer->layerId();
  1758. if (!isBatchedAtLevel(tensors, cur_level)) {
  1759. return at::_ops::atleast_1d_Sequence::call(tensors);
  1760. }
  1761. auto results = batch_rule(tensors);
  1762. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  1763. }
  1764. template <typename batch_rule_t, batch_rule_t batch_rule>
  1765. at::Tensor atleast_2d_generated_plumbing(const at::Tensor & self) {
  1766. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1767. auto maybe_layer = maybeCurrentDynamicLayer();
  1768. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1769. int64_t cur_level = maybe_layer->layerId();
  1770. if (!isBatchedAtLevel(self, cur_level)) {
  1771. return at::_ops::atleast_2d::call(self);
  1772. }
  1773. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1774. auto results = batch_rule(self_value, self_bdim);
  1775. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1776. }
  1777. template <typename batch_rule_t, batch_rule_t batch_rule>
  1778. ::std::vector<at::Tensor> atleast_2d_Sequence_generated_plumbing(at::TensorList tensors) {
  1779. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1780. auto maybe_layer = maybeCurrentDynamicLayer();
  1781. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1782. int64_t cur_level = maybe_layer->layerId();
  1783. if (!isBatchedAtLevel(tensors, cur_level)) {
  1784. return at::_ops::atleast_2d_Sequence::call(tensors);
  1785. }
  1786. auto results = batch_rule(tensors);
  1787. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  1788. }
  1789. template <typename batch_rule_t, batch_rule_t batch_rule>
  1790. at::Tensor atleast_3d_generated_plumbing(const at::Tensor & self) {
  1791. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1792. auto maybe_layer = maybeCurrentDynamicLayer();
  1793. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1794. int64_t cur_level = maybe_layer->layerId();
  1795. if (!isBatchedAtLevel(self, cur_level)) {
  1796. return at::_ops::atleast_3d::call(self);
  1797. }
  1798. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1799. auto results = batch_rule(self_value, self_bdim);
  1800. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1801. }
  1802. template <typename batch_rule_t, batch_rule_t batch_rule>
  1803. ::std::vector<at::Tensor> atleast_3d_Sequence_generated_plumbing(at::TensorList tensors) {
  1804. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1805. auto maybe_layer = maybeCurrentDynamicLayer();
  1806. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1807. int64_t cur_level = maybe_layer->layerId();
  1808. if (!isBatchedAtLevel(tensors, cur_level)) {
  1809. return at::_ops::atleast_3d_Sequence::call(tensors);
  1810. }
  1811. auto results = batch_rule(tensors);
  1812. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  1813. }
  1814. template <typename batch_rule_t, batch_rule_t batch_rule>
  1815. at::Tensor baddbmm_generated_plumbing(const at::Tensor & self, const at::Tensor & batch1, const at::Tensor & batch2, const at::Scalar & beta, const at::Scalar & alpha) {
  1816. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1817. auto maybe_layer = maybeCurrentDynamicLayer();
  1818. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1819. int64_t cur_level = maybe_layer->layerId();
  1820. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(batch1, cur_level) && !isBatchedAtLevel(batch2, cur_level)) {
  1821. return at::_ops::baddbmm::call(self, batch1, batch2, beta, alpha);
  1822. }
  1823. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1824. auto [batch1_value, batch1_bdim] = unwrapTensorAtLevel(batch1, cur_level);
  1825. auto [batch2_value, batch2_bdim] = unwrapTensorAtLevel(batch2, cur_level);
  1826. auto results = batch_rule(self_value, self_bdim, batch1_value, batch1_bdim, batch2_value, batch2_bdim, beta, alpha);
  1827. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1828. }
  1829. template <typename batch_rule_t, batch_rule_t batch_rule>
  1830. at::Tensor & baddbmm__generated_plumbing(at::Tensor & self, const at::Tensor & batch1, const at::Tensor & batch2, const at::Scalar & beta, const at::Scalar & alpha) {
  1831. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1832. auto maybe_layer = maybeCurrentDynamicLayer();
  1833. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  1834. int64_t cur_level = maybe_layer->layerId();
  1835. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(batch1, cur_level) && !isBatchedAtLevel(batch2, cur_level)) {
  1836. return at::_ops::baddbmm_::call(self, batch1, batch2, beta, alpha);
  1837. }
  1838. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1839. auto [batch1_value, batch1_bdim] = unwrapTensorAtLevel(batch1, cur_level);
  1840. auto [batch2_value, batch2_bdim] = unwrapTensorAtLevel(batch2, cur_level);
  1841. batch_rule(self_value, self_bdim, batch1_value, batch1_bdim, batch2_value, batch2_bdim, beta, alpha);
  1842. return self;
  1843. }
  1844. template <typename batch_rule_t, batch_rule_t batch_rule>
  1845. at::Tensor baddbmm_dtype_generated_plumbing(const at::Tensor & self, const at::Tensor & batch1, const at::Tensor & batch2, at::ScalarType out_dtype, const at::Scalar & beta, const at::Scalar & alpha) {
  1846. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1847. auto maybe_layer = maybeCurrentDynamicLayer();
  1848. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1849. int64_t cur_level = maybe_layer->layerId();
  1850. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(batch1, cur_level) && !isBatchedAtLevel(batch2, cur_level)) {
  1851. return at::_ops::baddbmm_dtype::call(self, batch1, batch2, out_dtype, beta, alpha);
  1852. }
  1853. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1854. auto [batch1_value, batch1_bdim] = unwrapTensorAtLevel(batch1, cur_level);
  1855. auto [batch2_value, batch2_bdim] = unwrapTensorAtLevel(batch2, cur_level);
  1856. auto results = batch_rule(self_value, self_bdim, batch1_value, batch1_bdim, batch2_value, batch2_bdim, out_dtype, beta, alpha);
  1857. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1858. }
  1859. template <typename batch_rule_t, batch_rule_t batch_rule>
  1860. at::Tensor batch_norm_generated_plumbing(const at::Tensor & input, const ::std::optional<at::Tensor> & weight, const ::std::optional<at::Tensor> & bias, const ::std::optional<at::Tensor> & running_mean, const ::std::optional<at::Tensor> & running_var, bool training, double momentum, double eps, bool cudnn_enabled) {
  1861. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1862. auto maybe_layer = maybeCurrentDynamicLayer();
  1863. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1864. int64_t cur_level = maybe_layer->layerId();
  1865. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level) && !isBatchedAtLevel(running_mean, cur_level) && !isBatchedAtLevel(running_var, cur_level)) {
  1866. return at::_ops::batch_norm::call(input, weight, bias, running_mean, running_var, training, momentum, eps, cudnn_enabled);
  1867. }
  1868. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  1869. std::optional<Tensor> weight_value;
  1870. std::optional<int64_t> weight_bdim;
  1871. if (weight) {
  1872. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  1873. }
  1874. std::optional<Tensor> bias_value;
  1875. std::optional<int64_t> bias_bdim;
  1876. if (bias) {
  1877. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  1878. }
  1879. std::optional<Tensor> running_mean_value;
  1880. std::optional<int64_t> running_mean_bdim;
  1881. if (running_mean) {
  1882. std::tie(running_mean_value, running_mean_bdim) = unwrapTensorAtLevel(running_mean.value(), cur_level);
  1883. }
  1884. std::optional<Tensor> running_var_value;
  1885. std::optional<int64_t> running_var_bdim;
  1886. if (running_var) {
  1887. std::tie(running_var_value, running_var_bdim) = unwrapTensorAtLevel(running_var.value(), cur_level);
  1888. }
  1889. auto results = batch_rule(input_value, input_bdim, weight_value, weight_bdim, bias_value, bias_bdim, running_mean_value, running_mean_bdim, running_var_value, running_var_bdim, training, momentum, eps, cudnn_enabled);
  1890. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1891. }
  1892. template <typename batch_rule_t, batch_rule_t batch_rule>
  1893. at::Tensor quantized_batch_norm_generated_plumbing(const at::Tensor & input, const ::std::optional<at::Tensor> & weight, const ::std::optional<at::Tensor> & bias, const at::Tensor & mean, const at::Tensor & var, double eps, double output_scale, int64_t output_zero_point) {
  1894. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1895. auto maybe_layer = maybeCurrentDynamicLayer();
  1896. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1897. int64_t cur_level = maybe_layer->layerId();
  1898. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level) && !isBatchedAtLevel(mean, cur_level) && !isBatchedAtLevel(var, cur_level)) {
  1899. return at::_ops::quantized_batch_norm::call(input, weight, bias, mean, var, eps, output_scale, output_zero_point);
  1900. }
  1901. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  1902. auto [mean_value, mean_bdim] = unwrapTensorAtLevel(mean, cur_level);
  1903. auto [var_value, var_bdim] = unwrapTensorAtLevel(var, cur_level);
  1904. std::optional<Tensor> weight_value;
  1905. std::optional<int64_t> weight_bdim;
  1906. if (weight) {
  1907. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  1908. }
  1909. std::optional<Tensor> bias_value;
  1910. std::optional<int64_t> bias_bdim;
  1911. if (bias) {
  1912. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  1913. }
  1914. auto results = batch_rule(input_value, input_bdim, weight_value, weight_bdim, bias_value, bias_bdim, mean_value, mean_bdim, var_value, var_bdim, eps, output_scale, output_zero_point);
  1915. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1916. }
  1917. template <typename batch_rule_t, batch_rule_t batch_rule>
  1918. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> _batch_norm_impl_index_backward_generated_plumbing(int64_t impl_index, const at::Tensor & input, const at::Tensor & grad_output, const ::std::optional<at::Tensor> & weight, const ::std::optional<at::Tensor> & running_mean, const ::std::optional<at::Tensor> & running_var, const ::std::optional<at::Tensor> & save_mean, const ::std::optional<at::Tensor> & save_var_transform, bool train, double eps, ::std::array<bool,3> output_mask, const at::Tensor & reservedSpace) {
  1919. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1920. auto maybe_layer = maybeCurrentDynamicLayer();
  1921. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1922. int64_t cur_level = maybe_layer->layerId();
  1923. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(running_mean, cur_level) && !isBatchedAtLevel(running_var, cur_level) && !isBatchedAtLevel(save_mean, cur_level) && !isBatchedAtLevel(save_var_transform, cur_level) && !isBatchedAtLevel(reservedSpace, cur_level)) {
  1924. return at::_ops::_batch_norm_impl_index_backward::call(impl_index, input, grad_output, weight, running_mean, running_var, save_mean, save_var_transform, train, eps, output_mask, reservedSpace);
  1925. }
  1926. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  1927. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  1928. auto [reservedSpace_value, reservedSpace_bdim] = unwrapTensorAtLevel(reservedSpace, cur_level);
  1929. std::optional<Tensor> weight_value;
  1930. std::optional<int64_t> weight_bdim;
  1931. if (weight) {
  1932. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  1933. }
  1934. std::optional<Tensor> running_mean_value;
  1935. std::optional<int64_t> running_mean_bdim;
  1936. if (running_mean) {
  1937. std::tie(running_mean_value, running_mean_bdim) = unwrapTensorAtLevel(running_mean.value(), cur_level);
  1938. }
  1939. std::optional<Tensor> running_var_value;
  1940. std::optional<int64_t> running_var_bdim;
  1941. if (running_var) {
  1942. std::tie(running_var_value, running_var_bdim) = unwrapTensorAtLevel(running_var.value(), cur_level);
  1943. }
  1944. std::optional<Tensor> save_mean_value;
  1945. std::optional<int64_t> save_mean_bdim;
  1946. if (save_mean) {
  1947. std::tie(save_mean_value, save_mean_bdim) = unwrapTensorAtLevel(save_mean.value(), cur_level);
  1948. }
  1949. std::optional<Tensor> save_var_transform_value;
  1950. std::optional<int64_t> save_var_transform_bdim;
  1951. if (save_var_transform) {
  1952. std::tie(save_var_transform_value, save_var_transform_bdim) = unwrapTensorAtLevel(save_var_transform.value(), cur_level);
  1953. }
  1954. auto results = batch_rule(impl_index, input_value, input_bdim, grad_output_value, grad_output_bdim, weight_value, weight_bdim, running_mean_value, running_mean_bdim, running_var_value, running_var_bdim, save_mean_value, save_mean_bdim, save_var_transform_value, save_var_transform_bdim, train, eps, output_mask, reservedSpace_value, reservedSpace_bdim);
  1955. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  1956. }
  1957. template <typename batch_rule_t, batch_rule_t batch_rule>
  1958. at::Tensor bernoulli_generated_plumbing(const at::Tensor & self, ::std::optional<at::Generator> generator) {
  1959. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1960. auto maybe_layer = maybeCurrentDynamicLayer();
  1961. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  1962. int64_t cur_level = maybe_layer->layerId();
  1963. if (!isBatchedAtLevel(self, cur_level)) {
  1964. return at::_ops::bernoulli::call(self, generator);
  1965. }
  1966. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1967. auto results = batch_rule(self_value, self_bdim, generator);
  1968. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  1969. }
  1970. template <typename batch_rule_t, batch_rule_t batch_rule>
  1971. at::Tensor & bernoulli__Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & p, ::std::optional<at::Generator> generator) {
  1972. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1973. auto maybe_layer = maybeCurrentDynamicLayer();
  1974. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  1975. int64_t cur_level = maybe_layer->layerId();
  1976. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(p, cur_level)) {
  1977. return at::_ops::bernoulli__Tensor::call(self, p, generator);
  1978. }
  1979. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1980. auto [p_value, p_bdim] = unwrapTensorAtLevel(p, cur_level);
  1981. batch_rule(self_value, self_bdim, p_value, p_bdim, generator);
  1982. return self;
  1983. }
  1984. template <typename batch_rule_t, batch_rule_t batch_rule>
  1985. at::Tensor & bernoulli__float_generated_plumbing(at::Tensor & self, double p, ::std::optional<at::Generator> generator) {
  1986. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  1987. auto maybe_layer = maybeCurrentDynamicLayer();
  1988. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  1989. int64_t cur_level = maybe_layer->layerId();
  1990. if (!isBatchedAtLevel(self, cur_level)) {
  1991. return at::_ops::bernoulli__float::call(self, p, generator);
  1992. }
  1993. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  1994. batch_rule(self_value, self_bdim, p, generator);
  1995. return self;
  1996. }
  1997. template <typename batch_rule_t, batch_rule_t batch_rule>
  1998. at::Tensor bernoulli_p_generated_plumbing(const at::Tensor & self, double p, ::std::optional<at::Generator> generator) {
  1999. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2000. auto maybe_layer = maybeCurrentDynamicLayer();
  2001. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2002. int64_t cur_level = maybe_layer->layerId();
  2003. if (!isBatchedAtLevel(self, cur_level)) {
  2004. return at::_ops::bernoulli_p::call(self, p, generator);
  2005. }
  2006. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2007. auto results = batch_rule(self_value, self_bdim, p, generator);
  2008. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2009. }
  2010. template <typename batch_rule_t, batch_rule_t batch_rule>
  2011. at::Tensor bilinear_generated_plumbing(const at::Tensor & input1, const at::Tensor & input2, const at::Tensor & weight, const ::std::optional<at::Tensor> & bias) {
  2012. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2013. auto maybe_layer = maybeCurrentDynamicLayer();
  2014. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2015. int64_t cur_level = maybe_layer->layerId();
  2016. if (!isBatchedAtLevel(input1, cur_level) && !isBatchedAtLevel(input2, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  2017. return at::_ops::bilinear::call(input1, input2, weight, bias);
  2018. }
  2019. auto [input1_value, input1_bdim] = unwrapTensorAtLevel(input1, cur_level);
  2020. auto [input2_value, input2_bdim] = unwrapTensorAtLevel(input2, cur_level);
  2021. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  2022. std::optional<Tensor> bias_value;
  2023. std::optional<int64_t> bias_bdim;
  2024. if (bias) {
  2025. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  2026. }
  2027. auto results = batch_rule(input1_value, input1_bdim, input2_value, input2_bdim, weight_value, weight_bdim, bias_value, bias_bdim);
  2028. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2029. }
  2030. template <typename batch_rule_t, batch_rule_t batch_rule>
  2031. at::Tensor binary_cross_entropy_generated_plumbing(const at::Tensor & self, const at::Tensor & target, const ::std::optional<at::Tensor> & weight, int64_t reduction) {
  2032. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2033. auto maybe_layer = maybeCurrentDynamicLayer();
  2034. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2035. int64_t cur_level = maybe_layer->layerId();
  2036. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(target, cur_level) && !isBatchedAtLevel(weight, cur_level)) {
  2037. return at::_ops::binary_cross_entropy::call(self, target, weight, reduction);
  2038. }
  2039. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2040. auto [target_value, target_bdim] = unwrapTensorAtLevel(target, cur_level);
  2041. std::optional<Tensor> weight_value;
  2042. std::optional<int64_t> weight_bdim;
  2043. if (weight) {
  2044. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  2045. }
  2046. auto results = batch_rule(self_value, self_bdim, target_value, target_bdim, weight_value, weight_bdim, reduction);
  2047. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2048. }
  2049. template <typename batch_rule_t, batch_rule_t batch_rule>
  2050. at::Tensor binary_cross_entropy_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self, const at::Tensor & target, const ::std::optional<at::Tensor> & weight, int64_t reduction) {
  2051. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2052. auto maybe_layer = maybeCurrentDynamicLayer();
  2053. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2054. int64_t cur_level = maybe_layer->layerId();
  2055. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(target, cur_level) && !isBatchedAtLevel(weight, cur_level)) {
  2056. return at::_ops::binary_cross_entropy_backward::call(grad_output, self, target, weight, reduction);
  2057. }
  2058. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  2059. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2060. auto [target_value, target_bdim] = unwrapTensorAtLevel(target, cur_level);
  2061. std::optional<Tensor> weight_value;
  2062. std::optional<int64_t> weight_bdim;
  2063. if (weight) {
  2064. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  2065. }
  2066. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim, target_value, target_bdim, weight_value, weight_bdim, reduction);
  2067. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2068. }
  2069. template <typename batch_rule_t, batch_rule_t batch_rule>
  2070. at::Tensor binary_cross_entropy_with_logits_generated_plumbing(const at::Tensor & self, const at::Tensor & target, const ::std::optional<at::Tensor> & weight, const ::std::optional<at::Tensor> & pos_weight, int64_t reduction) {
  2071. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2072. auto maybe_layer = maybeCurrentDynamicLayer();
  2073. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2074. int64_t cur_level = maybe_layer->layerId();
  2075. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(target, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(pos_weight, cur_level)) {
  2076. return at::_ops::binary_cross_entropy_with_logits::call(self, target, weight, pos_weight, reduction);
  2077. }
  2078. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2079. auto [target_value, target_bdim] = unwrapTensorAtLevel(target, cur_level);
  2080. std::optional<Tensor> weight_value;
  2081. std::optional<int64_t> weight_bdim;
  2082. if (weight) {
  2083. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  2084. }
  2085. std::optional<Tensor> pos_weight_value;
  2086. std::optional<int64_t> pos_weight_bdim;
  2087. if (pos_weight) {
  2088. std::tie(pos_weight_value, pos_weight_bdim) = unwrapTensorAtLevel(pos_weight.value(), cur_level);
  2089. }
  2090. auto results = batch_rule(self_value, self_bdim, target_value, target_bdim, weight_value, weight_bdim, pos_weight_value, pos_weight_bdim, reduction);
  2091. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2092. }
  2093. template <typename batch_rule_t, batch_rule_t batch_rule>
  2094. at::Tensor bincount_generated_plumbing(const at::Tensor & self, const ::std::optional<at::Tensor> & weights, c10::SymInt minlength) {
  2095. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2096. auto maybe_layer = maybeCurrentDynamicLayer();
  2097. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2098. int64_t cur_level = maybe_layer->layerId();
  2099. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(weights, cur_level)) {
  2100. return at::_ops::bincount::call(self, weights, minlength);
  2101. }
  2102. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2103. std::optional<Tensor> weights_value;
  2104. std::optional<int64_t> weights_bdim;
  2105. if (weights) {
  2106. std::tie(weights_value, weights_bdim) = unwrapTensorAtLevel(weights.value(), cur_level);
  2107. }
  2108. auto results = batch_rule(self_value, self_bdim, weights_value, weights_bdim, minlength);
  2109. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2110. }
  2111. template <typename batch_rule_t, batch_rule_t batch_rule>
  2112. at::Tensor bitwise_not_generated_plumbing(const at::Tensor & self) {
  2113. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2114. auto maybe_layer = maybeCurrentDynamicLayer();
  2115. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2116. int64_t cur_level = maybe_layer->layerId();
  2117. if (!isBatchedAtLevel(self, cur_level)) {
  2118. return at::_ops::bitwise_not::call(self);
  2119. }
  2120. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2121. auto results = batch_rule(self_value, self_bdim);
  2122. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2123. }
  2124. template <typename batch_rule_t, batch_rule_t batch_rule>
  2125. at::Tensor & bitwise_not__generated_plumbing(at::Tensor & self) {
  2126. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2127. auto maybe_layer = maybeCurrentDynamicLayer();
  2128. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  2129. int64_t cur_level = maybe_layer->layerId();
  2130. if (!isBatchedAtLevel(self, cur_level)) {
  2131. return at::_ops::bitwise_not_::call(self);
  2132. }
  2133. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2134. batch_rule(self_value, self_bdim);
  2135. return self;
  2136. }
  2137. template <typename batch_rule_t, batch_rule_t batch_rule>
  2138. at::Tensor copysign_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  2139. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2140. auto maybe_layer = maybeCurrentDynamicLayer();
  2141. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2142. int64_t cur_level = maybe_layer->layerId();
  2143. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  2144. return at::_ops::copysign_Tensor::call(self, other);
  2145. }
  2146. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2147. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  2148. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  2149. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2150. }
  2151. template <typename batch_rule_t, batch_rule_t batch_rule>
  2152. at::Tensor & copysign__Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  2153. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2154. auto maybe_layer = maybeCurrentDynamicLayer();
  2155. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  2156. int64_t cur_level = maybe_layer->layerId();
  2157. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  2158. return at::_ops::copysign__Tensor::call(self, other);
  2159. }
  2160. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2161. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  2162. batch_rule(self_value, self_bdim, other_value, other_bdim);
  2163. return self;
  2164. }
  2165. template <typename batch_rule_t, batch_rule_t batch_rule>
  2166. at::Tensor copysign_Scalar_generated_plumbing(const at::Tensor & self, const at::Scalar & other) {
  2167. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2168. auto maybe_layer = maybeCurrentDynamicLayer();
  2169. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2170. int64_t cur_level = maybe_layer->layerId();
  2171. if (!isBatchedAtLevel(self, cur_level)) {
  2172. return at::_ops::copysign_Scalar::call(self, other);
  2173. }
  2174. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2175. auto results = batch_rule(self_value, self_bdim, other);
  2176. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2177. }
  2178. template <typename batch_rule_t, batch_rule_t batch_rule>
  2179. at::Tensor & copysign__Scalar_generated_plumbing(at::Tensor & self, const at::Scalar & other) {
  2180. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2181. auto maybe_layer = maybeCurrentDynamicLayer();
  2182. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  2183. int64_t cur_level = maybe_layer->layerId();
  2184. if (!isBatchedAtLevel(self, cur_level)) {
  2185. return at::_ops::copysign__Scalar::call(self, other);
  2186. }
  2187. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2188. batch_rule(self_value, self_bdim, other);
  2189. return self;
  2190. }
  2191. template <typename batch_rule_t, batch_rule_t batch_rule>
  2192. at::Tensor _lazy_clone_generated_plumbing(const at::Tensor & self) {
  2193. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2194. auto maybe_layer = maybeCurrentDynamicLayer();
  2195. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2196. int64_t cur_level = maybe_layer->layerId();
  2197. if (!isBatchedAtLevel(self, cur_level)) {
  2198. return at::_ops::_lazy_clone::call(self);
  2199. }
  2200. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2201. auto results = batch_rule(self_value, self_bdim);
  2202. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2203. }
  2204. template <typename batch_rule_t, batch_rule_t batch_rule>
  2205. at::Tensor logical_not_generated_plumbing(const at::Tensor & self) {
  2206. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2207. auto maybe_layer = maybeCurrentDynamicLayer();
  2208. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2209. int64_t cur_level = maybe_layer->layerId();
  2210. if (!isBatchedAtLevel(self, cur_level)) {
  2211. return at::_ops::logical_not::call(self);
  2212. }
  2213. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2214. auto results = batch_rule(self_value, self_bdim);
  2215. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2216. }
  2217. template <typename batch_rule_t, batch_rule_t batch_rule>
  2218. at::Tensor & logical_not__generated_plumbing(at::Tensor & self) {
  2219. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2220. auto maybe_layer = maybeCurrentDynamicLayer();
  2221. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  2222. int64_t cur_level = maybe_layer->layerId();
  2223. if (!isBatchedAtLevel(self, cur_level)) {
  2224. return at::_ops::logical_not_::call(self);
  2225. }
  2226. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2227. batch_rule(self_value, self_bdim);
  2228. return self;
  2229. }
  2230. template <typename batch_rule_t, batch_rule_t batch_rule>
  2231. at::Tensor logical_xor_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  2232. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2233. auto maybe_layer = maybeCurrentDynamicLayer();
  2234. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2235. int64_t cur_level = maybe_layer->layerId();
  2236. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  2237. return at::_ops::logical_xor::call(self, other);
  2238. }
  2239. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2240. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  2241. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  2242. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2243. }
  2244. template <typename batch_rule_t, batch_rule_t batch_rule>
  2245. at::Tensor & logical_xor__generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  2246. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2247. auto maybe_layer = maybeCurrentDynamicLayer();
  2248. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  2249. int64_t cur_level = maybe_layer->layerId();
  2250. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  2251. return at::_ops::logical_xor_::call(self, other);
  2252. }
  2253. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2254. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  2255. batch_rule(self_value, self_bdim, other_value, other_bdim);
  2256. return self;
  2257. }
  2258. template <typename batch_rule_t, batch_rule_t batch_rule>
  2259. at::Tensor logical_and_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  2260. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2261. auto maybe_layer = maybeCurrentDynamicLayer();
  2262. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2263. int64_t cur_level = maybe_layer->layerId();
  2264. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  2265. return at::_ops::logical_and::call(self, other);
  2266. }
  2267. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2268. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  2269. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  2270. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2271. }
  2272. template <typename batch_rule_t, batch_rule_t batch_rule>
  2273. at::Tensor & logical_and__generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  2274. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2275. auto maybe_layer = maybeCurrentDynamicLayer();
  2276. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  2277. int64_t cur_level = maybe_layer->layerId();
  2278. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  2279. return at::_ops::logical_and_::call(self, other);
  2280. }
  2281. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2282. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  2283. batch_rule(self_value, self_bdim, other_value, other_bdim);
  2284. return self;
  2285. }
  2286. template <typename batch_rule_t, batch_rule_t batch_rule>
  2287. at::Tensor logical_or_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  2288. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2289. auto maybe_layer = maybeCurrentDynamicLayer();
  2290. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2291. int64_t cur_level = maybe_layer->layerId();
  2292. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  2293. return at::_ops::logical_or::call(self, other);
  2294. }
  2295. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2296. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  2297. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  2298. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2299. }
  2300. template <typename batch_rule_t, batch_rule_t batch_rule>
  2301. at::Tensor & logical_or__generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  2302. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2303. auto maybe_layer = maybeCurrentDynamicLayer();
  2304. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  2305. int64_t cur_level = maybe_layer->layerId();
  2306. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  2307. return at::_ops::logical_or_::call(self, other);
  2308. }
  2309. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2310. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  2311. batch_rule(self_value, self_bdim, other_value, other_bdim);
  2312. return self;
  2313. }
  2314. template <typename batch_rule_t, batch_rule_t batch_rule>
  2315. at::Tensor bmm_generated_plumbing(const at::Tensor & self, const at::Tensor & mat2) {
  2316. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2317. auto maybe_layer = maybeCurrentDynamicLayer();
  2318. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2319. int64_t cur_level = maybe_layer->layerId();
  2320. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(mat2, cur_level)) {
  2321. return at::_ops::bmm::call(self, mat2);
  2322. }
  2323. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2324. auto [mat2_value, mat2_bdim] = unwrapTensorAtLevel(mat2, cur_level);
  2325. auto results = batch_rule(self_value, self_bdim, mat2_value, mat2_bdim);
  2326. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2327. }
  2328. template <typename batch_rule_t, batch_rule_t batch_rule>
  2329. at::Tensor bmm_dtype_generated_plumbing(const at::Tensor & self, const at::Tensor & mat2, at::ScalarType out_dtype) {
  2330. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2331. auto maybe_layer = maybeCurrentDynamicLayer();
  2332. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2333. int64_t cur_level = maybe_layer->layerId();
  2334. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(mat2, cur_level)) {
  2335. return at::_ops::bmm_dtype::call(self, mat2, out_dtype);
  2336. }
  2337. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2338. auto [mat2_value, mat2_bdim] = unwrapTensorAtLevel(mat2, cur_level);
  2339. auto results = batch_rule(self_value, self_bdim, mat2_value, mat2_bdim, out_dtype);
  2340. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2341. }
  2342. template <typename batch_rule_t, batch_rule_t batch_rule>
  2343. ::std::vector<at::Tensor> broadcast_tensors_generated_plumbing(at::TensorList tensors) {
  2344. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2345. auto maybe_layer = maybeCurrentDynamicLayer();
  2346. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2347. int64_t cur_level = maybe_layer->layerId();
  2348. if (!isBatchedAtLevel(tensors, cur_level)) {
  2349. return at::_ops::broadcast_tensors::call(tensors);
  2350. }
  2351. auto results = batch_rule(tensors);
  2352. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  2353. }
  2354. template <typename batch_rule_t, batch_rule_t batch_rule>
  2355. at::Tensor broadcast_to_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef size) {
  2356. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2357. auto maybe_layer = maybeCurrentDynamicLayer();
  2358. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2359. int64_t cur_level = maybe_layer->layerId();
  2360. if (!isBatchedAtLevel(self, cur_level)) {
  2361. return at::_ops::broadcast_to::call(self, size);
  2362. }
  2363. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2364. auto results = batch_rule(self_value, self_bdim, size);
  2365. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2366. }
  2367. template <typename batch_rule_t, batch_rule_t batch_rule>
  2368. at::Tensor _sparse_broadcast_to_generated_plumbing(const at::Tensor & self, at::IntArrayRef size) {
  2369. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2370. auto maybe_layer = maybeCurrentDynamicLayer();
  2371. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2372. int64_t cur_level = maybe_layer->layerId();
  2373. if (!isBatchedAtLevel(self, cur_level)) {
  2374. return at::_ops::_sparse_broadcast_to::call(self, size);
  2375. }
  2376. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2377. auto results = batch_rule(self_value, self_bdim, size);
  2378. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2379. }
  2380. template <typename batch_rule_t, batch_rule_t batch_rule>
  2381. at::Tensor cat_generated_plumbing(const at::ITensorListRef & tensors, int64_t dim) {
  2382. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2383. auto maybe_layer = maybeCurrentDynamicLayer();
  2384. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2385. int64_t cur_level = maybe_layer->layerId();
  2386. if (!isBatchedAtLevel(tensors, cur_level)) {
  2387. return at::_ops::cat::call(tensors, dim);
  2388. }
  2389. auto results = batch_rule(tensors, dim);
  2390. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2391. }
  2392. template <typename batch_rule_t, batch_rule_t batch_rule>
  2393. at::Tensor cat_names_generated_plumbing(at::TensorList tensors, at::Dimname dim) {
  2394. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2395. auto maybe_layer = maybeCurrentDynamicLayer();
  2396. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2397. int64_t cur_level = maybe_layer->layerId();
  2398. if (!isBatchedAtLevel(tensors, cur_level)) {
  2399. return at::_ops::cat_names::call(tensors, dim);
  2400. }
  2401. auto results = batch_rule(tensors, dim);
  2402. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2403. }
  2404. template <typename batch_rule_t, batch_rule_t batch_rule>
  2405. at::Tensor concat_generated_plumbing(at::TensorList tensors, int64_t dim) {
  2406. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2407. auto maybe_layer = maybeCurrentDynamicLayer();
  2408. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2409. int64_t cur_level = maybe_layer->layerId();
  2410. if (!isBatchedAtLevel(tensors, cur_level)) {
  2411. return at::_ops::concat::call(tensors, dim);
  2412. }
  2413. auto results = batch_rule(tensors, dim);
  2414. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2415. }
  2416. template <typename batch_rule_t, batch_rule_t batch_rule>
  2417. at::Tensor concat_names_generated_plumbing(at::TensorList tensors, at::Dimname dim) {
  2418. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2419. auto maybe_layer = maybeCurrentDynamicLayer();
  2420. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2421. int64_t cur_level = maybe_layer->layerId();
  2422. if (!isBatchedAtLevel(tensors, cur_level)) {
  2423. return at::_ops::concat_names::call(tensors, dim);
  2424. }
  2425. auto results = batch_rule(tensors, dim);
  2426. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2427. }
  2428. template <typename batch_rule_t, batch_rule_t batch_rule>
  2429. at::Tensor concatenate_generated_plumbing(at::TensorList tensors, int64_t dim) {
  2430. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2431. auto maybe_layer = maybeCurrentDynamicLayer();
  2432. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2433. int64_t cur_level = maybe_layer->layerId();
  2434. if (!isBatchedAtLevel(tensors, cur_level)) {
  2435. return at::_ops::concatenate::call(tensors, dim);
  2436. }
  2437. auto results = batch_rule(tensors, dim);
  2438. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2439. }
  2440. template <typename batch_rule_t, batch_rule_t batch_rule>
  2441. at::Tensor concatenate_names_generated_plumbing(at::TensorList tensors, at::Dimname dim) {
  2442. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2443. auto maybe_layer = maybeCurrentDynamicLayer();
  2444. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2445. int64_t cur_level = maybe_layer->layerId();
  2446. if (!isBatchedAtLevel(tensors, cur_level)) {
  2447. return at::_ops::concatenate_names::call(tensors, dim);
  2448. }
  2449. auto results = batch_rule(tensors, dim);
  2450. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2451. }
  2452. template <typename batch_rule_t, batch_rule_t batch_rule>
  2453. at::Tensor block_diag_generated_plumbing(at::TensorList tensors) {
  2454. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2455. auto maybe_layer = maybeCurrentDynamicLayer();
  2456. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2457. int64_t cur_level = maybe_layer->layerId();
  2458. if (!isBatchedAtLevel(tensors, cur_level)) {
  2459. return at::_ops::block_diag::call(tensors);
  2460. }
  2461. auto results = batch_rule(tensors);
  2462. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2463. }
  2464. template <typename batch_rule_t, batch_rule_t batch_rule>
  2465. at::Tensor ceil_generated_plumbing(const at::Tensor & self) {
  2466. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2467. auto maybe_layer = maybeCurrentDynamicLayer();
  2468. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2469. int64_t cur_level = maybe_layer->layerId();
  2470. if (!isBatchedAtLevel(self, cur_level)) {
  2471. return at::_ops::ceil::call(self);
  2472. }
  2473. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2474. auto results = batch_rule(self_value, self_bdim);
  2475. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2476. }
  2477. template <typename batch_rule_t, batch_rule_t batch_rule>
  2478. at::Tensor & ceil__generated_plumbing(at::Tensor & self) {
  2479. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2480. auto maybe_layer = maybeCurrentDynamicLayer();
  2481. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  2482. int64_t cur_level = maybe_layer->layerId();
  2483. if (!isBatchedAtLevel(self, cur_level)) {
  2484. return at::_ops::ceil_::call(self);
  2485. }
  2486. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2487. batch_rule(self_value, self_bdim);
  2488. return self;
  2489. }
  2490. template <typename batch_rule_t, batch_rule_t batch_rule>
  2491. at::Tensor chain_matmul_generated_plumbing(at::TensorList matrices) {
  2492. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2493. auto maybe_layer = maybeCurrentDynamicLayer();
  2494. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2495. int64_t cur_level = maybe_layer->layerId();
  2496. if (!isBatchedAtLevel(matrices, cur_level)) {
  2497. return at::_ops::chain_matmul::call(matrices);
  2498. }
  2499. auto results = batch_rule(matrices);
  2500. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2501. }
  2502. template <typename batch_rule_t, batch_rule_t batch_rule>
  2503. ::std::vector<at::Tensor> unsafe_chunk_generated_plumbing(const at::Tensor & self, int64_t chunks, int64_t dim) {
  2504. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2505. auto maybe_layer = maybeCurrentDynamicLayer();
  2506. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2507. int64_t cur_level = maybe_layer->layerId();
  2508. if (!isBatchedAtLevel(self, cur_level)) {
  2509. return at::_ops::unsafe_chunk::call(self, chunks, dim);
  2510. }
  2511. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2512. auto results = batch_rule(self_value, self_bdim, chunks, dim);
  2513. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  2514. }
  2515. template <typename batch_rule_t, batch_rule_t batch_rule>
  2516. ::std::vector<at::Tensor> chunk_generated_plumbing(const at::Tensor & self, int64_t chunks, int64_t dim) {
  2517. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2518. auto maybe_layer = maybeCurrentDynamicLayer();
  2519. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2520. int64_t cur_level = maybe_layer->layerId();
  2521. if (!isBatchedAtLevel(self, cur_level)) {
  2522. return at::_ops::chunk::call(self, chunks, dim);
  2523. }
  2524. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2525. auto results = batch_rule(self_value, self_bdim, chunks, dim);
  2526. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  2527. }
  2528. template <typename batch_rule_t, batch_rule_t batch_rule>
  2529. ::std::vector<at::Tensor> tensor_split_sections_generated_plumbing(const at::Tensor & self, c10::SymInt sections, int64_t dim) {
  2530. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2531. auto maybe_layer = maybeCurrentDynamicLayer();
  2532. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2533. int64_t cur_level = maybe_layer->layerId();
  2534. if (!isBatchedAtLevel(self, cur_level)) {
  2535. return at::_ops::tensor_split_sections::call(self, sections, dim);
  2536. }
  2537. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2538. auto results = batch_rule(self_value, self_bdim, sections, dim);
  2539. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  2540. }
  2541. template <typename batch_rule_t, batch_rule_t batch_rule>
  2542. ::std::vector<at::Tensor> tensor_split_indices_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef indices, int64_t dim) {
  2543. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2544. auto maybe_layer = maybeCurrentDynamicLayer();
  2545. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2546. int64_t cur_level = maybe_layer->layerId();
  2547. if (!isBatchedAtLevel(self, cur_level)) {
  2548. return at::_ops::tensor_split_indices::call(self, indices, dim);
  2549. }
  2550. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2551. auto results = batch_rule(self_value, self_bdim, indices, dim);
  2552. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  2553. }
  2554. template <typename batch_rule_t, batch_rule_t batch_rule>
  2555. ::std::vector<at::Tensor> tensor_split_tensor_indices_or_sections_generated_plumbing(const at::Tensor & self, const at::Tensor & tensor_indices_or_sections, int64_t dim) {
  2556. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2557. auto maybe_layer = maybeCurrentDynamicLayer();
  2558. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2559. int64_t cur_level = maybe_layer->layerId();
  2560. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(tensor_indices_or_sections, cur_level)) {
  2561. return at::_ops::tensor_split_tensor_indices_or_sections::call(self, tensor_indices_or_sections, dim);
  2562. }
  2563. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2564. auto [tensor_indices_or_sections_value, tensor_indices_or_sections_bdim] = unwrapTensorAtLevel(tensor_indices_or_sections, cur_level);
  2565. auto results = batch_rule(self_value, self_bdim, tensor_indices_or_sections_value, tensor_indices_or_sections_bdim, dim);
  2566. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  2567. }
  2568. template <typename batch_rule_t, batch_rule_t batch_rule>
  2569. at::Tensor clamp_generated_plumbing(const at::Tensor & self, const ::std::optional<at::Scalar> & min, const ::std::optional<at::Scalar> & max) {
  2570. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2571. auto maybe_layer = maybeCurrentDynamicLayer();
  2572. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2573. int64_t cur_level = maybe_layer->layerId();
  2574. if (!isBatchedAtLevel(self, cur_level)) {
  2575. return at::_ops::clamp::call(self, min, max);
  2576. }
  2577. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2578. auto results = batch_rule(self_value, self_bdim, min, max);
  2579. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2580. }
  2581. template <typename batch_rule_t, batch_rule_t batch_rule>
  2582. at::Tensor clamp_Tensor_generated_plumbing(const at::Tensor & self, const ::std::optional<at::Tensor> & min, const ::std::optional<at::Tensor> & max) {
  2583. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2584. auto maybe_layer = maybeCurrentDynamicLayer();
  2585. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2586. int64_t cur_level = maybe_layer->layerId();
  2587. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(min, cur_level) && !isBatchedAtLevel(max, cur_level)) {
  2588. return at::_ops::clamp_Tensor::call(self, min, max);
  2589. }
  2590. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2591. std::optional<Tensor> min_value;
  2592. std::optional<int64_t> min_bdim;
  2593. if (min) {
  2594. std::tie(min_value, min_bdim) = unwrapTensorAtLevel(min.value(), cur_level);
  2595. }
  2596. std::optional<Tensor> max_value;
  2597. std::optional<int64_t> max_bdim;
  2598. if (max) {
  2599. std::tie(max_value, max_bdim) = unwrapTensorAtLevel(max.value(), cur_level);
  2600. }
  2601. auto results = batch_rule(self_value, self_bdim, min_value, min_bdim, max_value, max_bdim);
  2602. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2603. }
  2604. template <typename batch_rule_t, batch_rule_t batch_rule>
  2605. at::Tensor & clamp__generated_plumbing(at::Tensor & self, const ::std::optional<at::Scalar> & min, const ::std::optional<at::Scalar> & max) {
  2606. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2607. auto maybe_layer = maybeCurrentDynamicLayer();
  2608. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  2609. int64_t cur_level = maybe_layer->layerId();
  2610. if (!isBatchedAtLevel(self, cur_level)) {
  2611. return at::_ops::clamp_::call(self, min, max);
  2612. }
  2613. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2614. batch_rule(self_value, self_bdim, min, max);
  2615. return self;
  2616. }
  2617. template <typename batch_rule_t, batch_rule_t batch_rule>
  2618. at::Tensor & clamp__Tensor_generated_plumbing(at::Tensor & self, const ::std::optional<at::Tensor> & min, const ::std::optional<at::Tensor> & max) {
  2619. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2620. auto maybe_layer = maybeCurrentDynamicLayer();
  2621. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  2622. int64_t cur_level = maybe_layer->layerId();
  2623. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(min, cur_level) && !isBatchedAtLevel(max, cur_level)) {
  2624. return at::_ops::clamp__Tensor::call(self, min, max);
  2625. }
  2626. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2627. std::optional<Tensor> min_value;
  2628. std::optional<int64_t> min_bdim;
  2629. if (min) {
  2630. std::tie(min_value, min_bdim) = unwrapTensorAtLevel(min.value(), cur_level);
  2631. }
  2632. std::optional<Tensor> max_value;
  2633. std::optional<int64_t> max_bdim;
  2634. if (max) {
  2635. std::tie(max_value, max_bdim) = unwrapTensorAtLevel(max.value(), cur_level);
  2636. }
  2637. batch_rule(self_value, self_bdim, min_value, min_bdim, max_value, max_bdim);
  2638. return self;
  2639. }
  2640. template <typename batch_rule_t, batch_rule_t batch_rule>
  2641. at::Tensor clamp_max_generated_plumbing(const at::Tensor & self, const at::Scalar & max) {
  2642. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2643. auto maybe_layer = maybeCurrentDynamicLayer();
  2644. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2645. int64_t cur_level = maybe_layer->layerId();
  2646. if (!isBatchedAtLevel(self, cur_level)) {
  2647. return at::_ops::clamp_max::call(self, max);
  2648. }
  2649. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2650. auto results = batch_rule(self_value, self_bdim, max);
  2651. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2652. }
  2653. template <typename batch_rule_t, batch_rule_t batch_rule>
  2654. at::Tensor clamp_max_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & max) {
  2655. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2656. auto maybe_layer = maybeCurrentDynamicLayer();
  2657. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2658. int64_t cur_level = maybe_layer->layerId();
  2659. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(max, cur_level)) {
  2660. return at::_ops::clamp_max_Tensor::call(self, max);
  2661. }
  2662. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2663. auto [max_value, max_bdim] = unwrapTensorAtLevel(max, cur_level);
  2664. auto results = batch_rule(self_value, self_bdim, max_value, max_bdim);
  2665. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2666. }
  2667. template <typename batch_rule_t, batch_rule_t batch_rule>
  2668. at::Tensor & clamp_max__generated_plumbing(at::Tensor & self, const at::Scalar & max) {
  2669. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2670. auto maybe_layer = maybeCurrentDynamicLayer();
  2671. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  2672. int64_t cur_level = maybe_layer->layerId();
  2673. if (!isBatchedAtLevel(self, cur_level)) {
  2674. return at::_ops::clamp_max_::call(self, max);
  2675. }
  2676. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2677. batch_rule(self_value, self_bdim, max);
  2678. return self;
  2679. }
  2680. template <typename batch_rule_t, batch_rule_t batch_rule>
  2681. at::Tensor & clamp_max__Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & max) {
  2682. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2683. auto maybe_layer = maybeCurrentDynamicLayer();
  2684. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  2685. int64_t cur_level = maybe_layer->layerId();
  2686. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(max, cur_level)) {
  2687. return at::_ops::clamp_max__Tensor::call(self, max);
  2688. }
  2689. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2690. auto [max_value, max_bdim] = unwrapTensorAtLevel(max, cur_level);
  2691. batch_rule(self_value, self_bdim, max_value, max_bdim);
  2692. return self;
  2693. }
  2694. template <typename batch_rule_t, batch_rule_t batch_rule>
  2695. at::Tensor clamp_min_generated_plumbing(const at::Tensor & self, const at::Scalar & min) {
  2696. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2697. auto maybe_layer = maybeCurrentDynamicLayer();
  2698. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2699. int64_t cur_level = maybe_layer->layerId();
  2700. if (!isBatchedAtLevel(self, cur_level)) {
  2701. return at::_ops::clamp_min::call(self, min);
  2702. }
  2703. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2704. auto results = batch_rule(self_value, self_bdim, min);
  2705. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2706. }
  2707. template <typename batch_rule_t, batch_rule_t batch_rule>
  2708. at::Tensor clamp_min_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & min) {
  2709. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2710. auto maybe_layer = maybeCurrentDynamicLayer();
  2711. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2712. int64_t cur_level = maybe_layer->layerId();
  2713. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(min, cur_level)) {
  2714. return at::_ops::clamp_min_Tensor::call(self, min);
  2715. }
  2716. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2717. auto [min_value, min_bdim] = unwrapTensorAtLevel(min, cur_level);
  2718. auto results = batch_rule(self_value, self_bdim, min_value, min_bdim);
  2719. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2720. }
  2721. template <typename batch_rule_t, batch_rule_t batch_rule>
  2722. at::Tensor & clamp_min__generated_plumbing(at::Tensor & self, const at::Scalar & min) {
  2723. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2724. auto maybe_layer = maybeCurrentDynamicLayer();
  2725. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  2726. int64_t cur_level = maybe_layer->layerId();
  2727. if (!isBatchedAtLevel(self, cur_level)) {
  2728. return at::_ops::clamp_min_::call(self, min);
  2729. }
  2730. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2731. batch_rule(self_value, self_bdim, min);
  2732. return self;
  2733. }
  2734. template <typename batch_rule_t, batch_rule_t batch_rule>
  2735. at::Tensor & clamp_min__Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & min) {
  2736. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2737. auto maybe_layer = maybeCurrentDynamicLayer();
  2738. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  2739. int64_t cur_level = maybe_layer->layerId();
  2740. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(min, cur_level)) {
  2741. return at::_ops::clamp_min__Tensor::call(self, min);
  2742. }
  2743. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2744. auto [min_value, min_bdim] = unwrapTensorAtLevel(min, cur_level);
  2745. batch_rule(self_value, self_bdim, min_value, min_bdim);
  2746. return self;
  2747. }
  2748. template <typename batch_rule_t, batch_rule_t batch_rule>
  2749. at::Tensor clip_generated_plumbing(const at::Tensor & self, const ::std::optional<at::Scalar> & min, const ::std::optional<at::Scalar> & max) {
  2750. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2751. auto maybe_layer = maybeCurrentDynamicLayer();
  2752. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2753. int64_t cur_level = maybe_layer->layerId();
  2754. if (!isBatchedAtLevel(self, cur_level)) {
  2755. return at::_ops::clip::call(self, min, max);
  2756. }
  2757. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2758. auto results = batch_rule(self_value, self_bdim, min, max);
  2759. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2760. }
  2761. template <typename batch_rule_t, batch_rule_t batch_rule>
  2762. at::Tensor clip_Tensor_generated_plumbing(const at::Tensor & self, const ::std::optional<at::Tensor> & min, const ::std::optional<at::Tensor> & max) {
  2763. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2764. auto maybe_layer = maybeCurrentDynamicLayer();
  2765. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2766. int64_t cur_level = maybe_layer->layerId();
  2767. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(min, cur_level) && !isBatchedAtLevel(max, cur_level)) {
  2768. return at::_ops::clip_Tensor::call(self, min, max);
  2769. }
  2770. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2771. std::optional<Tensor> min_value;
  2772. std::optional<int64_t> min_bdim;
  2773. if (min) {
  2774. std::tie(min_value, min_bdim) = unwrapTensorAtLevel(min.value(), cur_level);
  2775. }
  2776. std::optional<Tensor> max_value;
  2777. std::optional<int64_t> max_bdim;
  2778. if (max) {
  2779. std::tie(max_value, max_bdim) = unwrapTensorAtLevel(max.value(), cur_level);
  2780. }
  2781. auto results = batch_rule(self_value, self_bdim, min_value, min_bdim, max_value, max_bdim);
  2782. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2783. }
  2784. template <typename batch_rule_t, batch_rule_t batch_rule>
  2785. at::Tensor & clip__generated_plumbing(at::Tensor & self, const ::std::optional<at::Scalar> & min, const ::std::optional<at::Scalar> & max) {
  2786. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2787. auto maybe_layer = maybeCurrentDynamicLayer();
  2788. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  2789. int64_t cur_level = maybe_layer->layerId();
  2790. if (!isBatchedAtLevel(self, cur_level)) {
  2791. return at::_ops::clip_::call(self, min, max);
  2792. }
  2793. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2794. batch_rule(self_value, self_bdim, min, max);
  2795. return self;
  2796. }
  2797. template <typename batch_rule_t, batch_rule_t batch_rule>
  2798. at::Tensor & clip__Tensor_generated_plumbing(at::Tensor & self, const ::std::optional<at::Tensor> & min, const ::std::optional<at::Tensor> & max) {
  2799. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2800. auto maybe_layer = maybeCurrentDynamicLayer();
  2801. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  2802. int64_t cur_level = maybe_layer->layerId();
  2803. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(min, cur_level) && !isBatchedAtLevel(max, cur_level)) {
  2804. return at::_ops::clip__Tensor::call(self, min, max);
  2805. }
  2806. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2807. std::optional<Tensor> min_value;
  2808. std::optional<int64_t> min_bdim;
  2809. if (min) {
  2810. std::tie(min_value, min_bdim) = unwrapTensorAtLevel(min.value(), cur_level);
  2811. }
  2812. std::optional<Tensor> max_value;
  2813. std::optional<int64_t> max_bdim;
  2814. if (max) {
  2815. std::tie(max_value, max_bdim) = unwrapTensorAtLevel(max.value(), cur_level);
  2816. }
  2817. batch_rule(self_value, self_bdim, min_value, min_bdim, max_value, max_bdim);
  2818. return self;
  2819. }
  2820. template <typename batch_rule_t, batch_rule_t batch_rule>
  2821. at::Tensor complex_generated_plumbing(const at::Tensor & real, const at::Tensor & imag) {
  2822. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2823. auto maybe_layer = maybeCurrentDynamicLayer();
  2824. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2825. int64_t cur_level = maybe_layer->layerId();
  2826. if (!isBatchedAtLevel(real, cur_level) && !isBatchedAtLevel(imag, cur_level)) {
  2827. return at::_ops::complex::call(real, imag);
  2828. }
  2829. auto [real_value, real_bdim] = unwrapTensorAtLevel(real, cur_level);
  2830. auto [imag_value, imag_bdim] = unwrapTensorAtLevel(imag, cur_level);
  2831. auto results = batch_rule(real_value, real_bdim, imag_value, imag_bdim);
  2832. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2833. }
  2834. template <typename batch_rule_t, batch_rule_t batch_rule>
  2835. at::Tensor polar_generated_plumbing(const at::Tensor & abs, const at::Tensor & angle) {
  2836. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2837. auto maybe_layer = maybeCurrentDynamicLayer();
  2838. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2839. int64_t cur_level = maybe_layer->layerId();
  2840. if (!isBatchedAtLevel(abs, cur_level) && !isBatchedAtLevel(angle, cur_level)) {
  2841. return at::_ops::polar::call(abs, angle);
  2842. }
  2843. auto [abs_value, abs_bdim] = unwrapTensorAtLevel(abs, cur_level);
  2844. auto [angle_value, angle_bdim] = unwrapTensorAtLevel(angle, cur_level);
  2845. auto results = batch_rule(abs_value, abs_bdim, angle_value, angle_bdim);
  2846. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2847. }
  2848. template <typename batch_rule_t, batch_rule_t batch_rule>
  2849. at::Tensor constant_pad_nd_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef pad, const at::Scalar & value) {
  2850. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2851. auto maybe_layer = maybeCurrentDynamicLayer();
  2852. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2853. int64_t cur_level = maybe_layer->layerId();
  2854. if (!isBatchedAtLevel(self, cur_level)) {
  2855. return at::_ops::constant_pad_nd::call(self, pad, value);
  2856. }
  2857. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2858. auto results = batch_rule(self_value, self_bdim, pad, value);
  2859. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2860. }
  2861. template <typename batch_rule_t, batch_rule_t batch_rule>
  2862. at::Tensor contiguous_generated_plumbing(const at::Tensor & self, at::MemoryFormat memory_format) {
  2863. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2864. auto maybe_layer = maybeCurrentDynamicLayer();
  2865. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2866. int64_t cur_level = maybe_layer->layerId();
  2867. if (!isBatchedAtLevel(self, cur_level)) {
  2868. return at::_ops::contiguous::call(self, memory_format);
  2869. }
  2870. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  2871. auto results = batch_rule(self_value, self_bdim, memory_format);
  2872. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2873. }
  2874. template <typename batch_rule_t, batch_rule_t batch_rule>
  2875. at::Tensor convolution_generated_plumbing(const at::Tensor & input, const at::Tensor & weight, const ::std::optional<at::Tensor> & bias, c10::SymIntArrayRef stride, c10::SymIntArrayRef padding, c10::SymIntArrayRef dilation, bool transposed, c10::SymIntArrayRef output_padding, c10::SymInt groups) {
  2876. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2877. auto maybe_layer = maybeCurrentDynamicLayer();
  2878. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2879. int64_t cur_level = maybe_layer->layerId();
  2880. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  2881. return at::_ops::convolution::call(input, weight, bias, stride, padding, dilation, transposed, output_padding, groups);
  2882. }
  2883. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  2884. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  2885. std::optional<Tensor> bias_value;
  2886. std::optional<int64_t> bias_bdim;
  2887. if (bias) {
  2888. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  2889. }
  2890. auto results = batch_rule(input_value, input_bdim, weight_value, weight_bdim, bias_value, bias_bdim, stride, padding, dilation, transposed, output_padding, groups);
  2891. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2892. }
  2893. template <typename batch_rule_t, batch_rule_t batch_rule>
  2894. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> convolution_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & input, const at::Tensor & weight, at::OptionalSymIntArrayRef bias_sizes, c10::SymIntArrayRef stride, c10::SymIntArrayRef padding, c10::SymIntArrayRef dilation, bool transposed, c10::SymIntArrayRef output_padding, c10::SymInt groups, ::std::array<bool,3> output_mask) {
  2895. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2896. auto maybe_layer = maybeCurrentDynamicLayer();
  2897. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2898. int64_t cur_level = maybe_layer->layerId();
  2899. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level)) {
  2900. return at::_ops::convolution_backward::call(grad_output, input, weight, bias_sizes, stride, padding, dilation, transposed, output_padding, groups, output_mask);
  2901. }
  2902. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  2903. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  2904. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  2905. auto results = batch_rule(grad_output_value, grad_output_bdim, input_value, input_bdim, weight_value, weight_bdim, bias_sizes, stride, padding, dilation, transposed, output_padding, groups, output_mask);
  2906. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  2907. }
  2908. template <typename batch_rule_t, batch_rule_t batch_rule>
  2909. at::Tensor convolution_overrideable_generated_plumbing(const at::Tensor & input, const at::Tensor & weight, const ::std::optional<at::Tensor> & bias, c10::SymIntArrayRef stride, c10::SymIntArrayRef padding, c10::SymIntArrayRef dilation, bool transposed, c10::SymIntArrayRef output_padding, c10::SymInt groups) {
  2910. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2911. auto maybe_layer = maybeCurrentDynamicLayer();
  2912. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2913. int64_t cur_level = maybe_layer->layerId();
  2914. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  2915. return at::_ops::convolution_overrideable::call(input, weight, bias, stride, padding, dilation, transposed, output_padding, groups);
  2916. }
  2917. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  2918. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  2919. std::optional<Tensor> bias_value;
  2920. std::optional<int64_t> bias_bdim;
  2921. if (bias) {
  2922. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  2923. }
  2924. auto results = batch_rule(input_value, input_bdim, weight_value, weight_bdim, bias_value, bias_bdim, stride, padding, dilation, transposed, output_padding, groups);
  2925. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2926. }
  2927. template <typename batch_rule_t, batch_rule_t batch_rule>
  2928. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> convolution_backward_overrideable_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & input, const at::Tensor & weight, c10::SymIntArrayRef stride, c10::SymIntArrayRef padding, c10::SymIntArrayRef dilation, bool transposed, c10::SymIntArrayRef output_padding, c10::SymInt groups, ::std::array<bool,3> output_mask) {
  2929. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2930. auto maybe_layer = maybeCurrentDynamicLayer();
  2931. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2932. int64_t cur_level = maybe_layer->layerId();
  2933. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level)) {
  2934. return at::_ops::convolution_backward_overrideable::call(grad_output, input, weight, stride, padding, dilation, transposed, output_padding, groups, output_mask);
  2935. }
  2936. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  2937. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  2938. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  2939. auto results = batch_rule(grad_output_value, grad_output_bdim, input_value, input_bdim, weight_value, weight_bdim, stride, padding, dilation, transposed, output_padding, groups, output_mask);
  2940. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  2941. }
  2942. template <typename batch_rule_t, batch_rule_t batch_rule>
  2943. at::Tensor _convolution_generated_plumbing(const at::Tensor & input, const at::Tensor & weight, const ::std::optional<at::Tensor> & bias, c10::SymIntArrayRef stride, c10::SymIntArrayRef padding, c10::SymIntArrayRef dilation, bool transposed, c10::SymIntArrayRef output_padding, c10::SymInt groups, bool benchmark, bool deterministic, bool cudnn_enabled, bool allow_tf32) {
  2944. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2945. auto maybe_layer = maybeCurrentDynamicLayer();
  2946. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2947. int64_t cur_level = maybe_layer->layerId();
  2948. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  2949. return at::_ops::_convolution::call(input, weight, bias, stride, padding, dilation, transposed, output_padding, groups, benchmark, deterministic, cudnn_enabled, allow_tf32);
  2950. }
  2951. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  2952. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  2953. std::optional<Tensor> bias_value;
  2954. std::optional<int64_t> bias_bdim;
  2955. if (bias) {
  2956. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  2957. }
  2958. auto results = batch_rule(input_value, input_bdim, weight_value, weight_bdim, bias_value, bias_bdim, stride, padding, dilation, transposed, output_padding, groups, benchmark, deterministic, cudnn_enabled, allow_tf32);
  2959. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2960. }
  2961. template <typename batch_rule_t, batch_rule_t batch_rule>
  2962. at::Tensor _convolution_deprecated_generated_plumbing(const at::Tensor & input, const at::Tensor & weight, const ::std::optional<at::Tensor> & bias, c10::SymIntArrayRef stride, c10::SymIntArrayRef padding, c10::SymIntArrayRef dilation, bool transposed, at::IntArrayRef output_padding, c10::SymInt groups, bool benchmark, bool deterministic, bool cudnn_enabled) {
  2963. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2964. auto maybe_layer = maybeCurrentDynamicLayer();
  2965. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2966. int64_t cur_level = maybe_layer->layerId();
  2967. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  2968. return at::_ops::_convolution_deprecated::call(input, weight, bias, stride, padding, dilation, transposed, output_padding, groups, benchmark, deterministic, cudnn_enabled);
  2969. }
  2970. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  2971. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  2972. std::optional<Tensor> bias_value;
  2973. std::optional<int64_t> bias_bdim;
  2974. if (bias) {
  2975. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  2976. }
  2977. auto results = batch_rule(input_value, input_bdim, weight_value, weight_bdim, bias_value, bias_bdim, stride, padding, dilation, transposed, output_padding, groups, benchmark, deterministic, cudnn_enabled);
  2978. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2979. }
  2980. template <typename batch_rule_t, batch_rule_t batch_rule>
  2981. at::Tensor _convolution_mode_generated_plumbing(const at::Tensor & input, const at::Tensor & weight, const ::std::optional<at::Tensor> & bias, c10::SymIntArrayRef stride, c10::string_view padding, c10::SymIntArrayRef dilation, c10::SymInt groups) {
  2982. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  2983. auto maybe_layer = maybeCurrentDynamicLayer();
  2984. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  2985. int64_t cur_level = maybe_layer->layerId();
  2986. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  2987. return at::_ops::_convolution_mode::call(input, weight, bias, stride, padding, dilation, groups);
  2988. }
  2989. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  2990. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  2991. std::optional<Tensor> bias_value;
  2992. std::optional<int64_t> bias_bdim;
  2993. if (bias) {
  2994. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  2995. }
  2996. auto results = batch_rule(input_value, input_bdim, weight_value, weight_bdim, bias_value, bias_bdim, stride, padding, dilation, groups);
  2997. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  2998. }
  2999. template <typename batch_rule_t, batch_rule_t batch_rule>
  3000. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> _convolution_double_backward_generated_plumbing(const ::std::optional<at::Tensor> & ggI, const ::std::optional<at::Tensor> & ggW, const ::std::optional<at::Tensor> & ggb, const at::Tensor & gO, const at::Tensor & weight, const at::Tensor & self, c10::SymIntArrayRef stride, c10::SymIntArrayRef padding, c10::SymIntArrayRef dilation, bool transposed, c10::SymIntArrayRef output_padding, c10::SymInt groups, ::std::array<bool,3> output_mask) {
  3001. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3002. auto maybe_layer = maybeCurrentDynamicLayer();
  3003. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3004. int64_t cur_level = maybe_layer->layerId();
  3005. if (!isBatchedAtLevel(ggI, cur_level) && !isBatchedAtLevel(ggW, cur_level) && !isBatchedAtLevel(ggb, cur_level) && !isBatchedAtLevel(gO, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(self, cur_level)) {
  3006. return at::_ops::_convolution_double_backward::call(ggI, ggW, ggb, gO, weight, self, stride, padding, dilation, transposed, output_padding, groups, output_mask);
  3007. }
  3008. auto [gO_value, gO_bdim] = unwrapTensorAtLevel(gO, cur_level);
  3009. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  3010. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  3011. std::optional<Tensor> ggI_value;
  3012. std::optional<int64_t> ggI_bdim;
  3013. if (ggI) {
  3014. std::tie(ggI_value, ggI_bdim) = unwrapTensorAtLevel(ggI.value(), cur_level);
  3015. }
  3016. std::optional<Tensor> ggW_value;
  3017. std::optional<int64_t> ggW_bdim;
  3018. if (ggW) {
  3019. std::tie(ggW_value, ggW_bdim) = unwrapTensorAtLevel(ggW.value(), cur_level);
  3020. }
  3021. std::optional<Tensor> ggb_value;
  3022. std::optional<int64_t> ggb_bdim;
  3023. if (ggb) {
  3024. std::tie(ggb_value, ggb_bdim) = unwrapTensorAtLevel(ggb.value(), cur_level);
  3025. }
  3026. auto results = batch_rule(ggI_value, ggI_bdim, ggW_value, ggW_bdim, ggb_value, ggb_bdim, gO_value, gO_bdim, weight_value, weight_bdim, self_value, self_bdim, stride, padding, dilation, transposed, output_padding, groups, output_mask);
  3027. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  3028. }
  3029. template <typename batch_rule_t, batch_rule_t batch_rule>
  3030. at::Tensor conv1d_generated_plumbing(const at::Tensor & input, const at::Tensor & weight, const ::std::optional<at::Tensor> & bias, c10::SymIntArrayRef stride, c10::SymIntArrayRef padding, c10::SymIntArrayRef dilation, c10::SymInt groups) {
  3031. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3032. auto maybe_layer = maybeCurrentDynamicLayer();
  3033. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3034. int64_t cur_level = maybe_layer->layerId();
  3035. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  3036. return at::_ops::conv1d::call(input, weight, bias, stride, padding, dilation, groups);
  3037. }
  3038. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  3039. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  3040. std::optional<Tensor> bias_value;
  3041. std::optional<int64_t> bias_bdim;
  3042. if (bias) {
  3043. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  3044. }
  3045. auto results = batch_rule(input_value, input_bdim, weight_value, weight_bdim, bias_value, bias_bdim, stride, padding, dilation, groups);
  3046. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  3047. }
  3048. template <typename batch_rule_t, batch_rule_t batch_rule>
  3049. at::Tensor conv2d_generated_plumbing(const at::Tensor & input, const at::Tensor & weight, const ::std::optional<at::Tensor> & bias, c10::SymIntArrayRef stride, c10::SymIntArrayRef padding, c10::SymIntArrayRef dilation, c10::SymInt groups) {
  3050. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3051. auto maybe_layer = maybeCurrentDynamicLayer();
  3052. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3053. int64_t cur_level = maybe_layer->layerId();
  3054. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  3055. return at::_ops::conv2d::call(input, weight, bias, stride, padding, dilation, groups);
  3056. }
  3057. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  3058. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  3059. std::optional<Tensor> bias_value;
  3060. std::optional<int64_t> bias_bdim;
  3061. if (bias) {
  3062. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  3063. }
  3064. auto results = batch_rule(input_value, input_bdim, weight_value, weight_bdim, bias_value, bias_bdim, stride, padding, dilation, groups);
  3065. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  3066. }
  3067. template <typename batch_rule_t, batch_rule_t batch_rule>
  3068. at::Tensor conv3d_generated_plumbing(const at::Tensor & input, const at::Tensor & weight, const ::std::optional<at::Tensor> & bias, c10::SymIntArrayRef stride, c10::SymIntArrayRef padding, c10::SymIntArrayRef dilation, c10::SymInt groups) {
  3069. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3070. auto maybe_layer = maybeCurrentDynamicLayer();
  3071. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3072. int64_t cur_level = maybe_layer->layerId();
  3073. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  3074. return at::_ops::conv3d::call(input, weight, bias, stride, padding, dilation, groups);
  3075. }
  3076. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  3077. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  3078. std::optional<Tensor> bias_value;
  3079. std::optional<int64_t> bias_bdim;
  3080. if (bias) {
  3081. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  3082. }
  3083. auto results = batch_rule(input_value, input_bdim, weight_value, weight_bdim, bias_value, bias_bdim, stride, padding, dilation, groups);
  3084. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  3085. }
  3086. template <typename batch_rule_t, batch_rule_t batch_rule>
  3087. at::Tensor conv1d_padding_generated_plumbing(const at::Tensor & input, const at::Tensor & weight, const ::std::optional<at::Tensor> & bias, c10::SymIntArrayRef stride, c10::string_view padding, c10::SymIntArrayRef dilation, c10::SymInt groups) {
  3088. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3089. auto maybe_layer = maybeCurrentDynamicLayer();
  3090. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3091. int64_t cur_level = maybe_layer->layerId();
  3092. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  3093. return at::_ops::conv1d_padding::call(input, weight, bias, stride, padding, dilation, groups);
  3094. }
  3095. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  3096. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  3097. std::optional<Tensor> bias_value;
  3098. std::optional<int64_t> bias_bdim;
  3099. if (bias) {
  3100. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  3101. }
  3102. auto results = batch_rule(input_value, input_bdim, weight_value, weight_bdim, bias_value, bias_bdim, stride, padding, dilation, groups);
  3103. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  3104. }
  3105. template <typename batch_rule_t, batch_rule_t batch_rule>
  3106. at::Tensor conv2d_padding_generated_plumbing(const at::Tensor & input, const at::Tensor & weight, const ::std::optional<at::Tensor> & bias, c10::SymIntArrayRef stride, c10::string_view padding, c10::SymIntArrayRef dilation, c10::SymInt groups) {
  3107. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3108. auto maybe_layer = maybeCurrentDynamicLayer();
  3109. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3110. int64_t cur_level = maybe_layer->layerId();
  3111. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  3112. return at::_ops::conv2d_padding::call(input, weight, bias, stride, padding, dilation, groups);
  3113. }
  3114. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  3115. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  3116. std::optional<Tensor> bias_value;
  3117. std::optional<int64_t> bias_bdim;
  3118. if (bias) {
  3119. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  3120. }
  3121. auto results = batch_rule(input_value, input_bdim, weight_value, weight_bdim, bias_value, bias_bdim, stride, padding, dilation, groups);
  3122. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  3123. }
  3124. template <typename batch_rule_t, batch_rule_t batch_rule>
  3125. at::Tensor conv3d_padding_generated_plumbing(const at::Tensor & input, const at::Tensor & weight, const ::std::optional<at::Tensor> & bias, c10::SymIntArrayRef stride, c10::string_view padding, c10::SymIntArrayRef dilation, c10::SymInt groups) {
  3126. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3127. auto maybe_layer = maybeCurrentDynamicLayer();
  3128. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3129. int64_t cur_level = maybe_layer->layerId();
  3130. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  3131. return at::_ops::conv3d_padding::call(input, weight, bias, stride, padding, dilation, groups);
  3132. }
  3133. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  3134. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  3135. std::optional<Tensor> bias_value;
  3136. std::optional<int64_t> bias_bdim;
  3137. if (bias) {
  3138. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  3139. }
  3140. auto results = batch_rule(input_value, input_bdim, weight_value, weight_bdim, bias_value, bias_bdim, stride, padding, dilation, groups);
  3141. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  3142. }
  3143. template <typename batch_rule_t, batch_rule_t batch_rule>
  3144. at::Tensor conv_tbc_generated_plumbing(const at::Tensor & self, const at::Tensor & weight, const at::Tensor & bias, int64_t pad) {
  3145. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3146. auto maybe_layer = maybeCurrentDynamicLayer();
  3147. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3148. int64_t cur_level = maybe_layer->layerId();
  3149. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  3150. return at::_ops::conv_tbc::call(self, weight, bias, pad);
  3151. }
  3152. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  3153. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  3154. auto [bias_value, bias_bdim] = unwrapTensorAtLevel(bias, cur_level);
  3155. auto results = batch_rule(self_value, self_bdim, weight_value, weight_bdim, bias_value, bias_bdim, pad);
  3156. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  3157. }
  3158. template <typename batch_rule_t, batch_rule_t batch_rule>
  3159. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> conv_tbc_backward_generated_plumbing(const at::Tensor & self, const at::Tensor & input, const at::Tensor & weight, const at::Tensor & bias, int64_t pad) {
  3160. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3161. auto maybe_layer = maybeCurrentDynamicLayer();
  3162. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3163. int64_t cur_level = maybe_layer->layerId();
  3164. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  3165. return at::_ops::conv_tbc_backward::call(self, input, weight, bias, pad);
  3166. }
  3167. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  3168. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  3169. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  3170. auto [bias_value, bias_bdim] = unwrapTensorAtLevel(bias, cur_level);
  3171. auto results = batch_rule(self_value, self_bdim, input_value, input_bdim, weight_value, weight_bdim, bias_value, bias_bdim, pad);
  3172. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  3173. }
  3174. template <typename batch_rule_t, batch_rule_t batch_rule>
  3175. at::Tensor conv_transpose1d_generated_plumbing(const at::Tensor & input, const at::Tensor & weight, const ::std::optional<at::Tensor> & bias, c10::SymIntArrayRef stride, c10::SymIntArrayRef padding, c10::SymIntArrayRef output_padding, c10::SymInt groups, c10::SymIntArrayRef dilation) {
  3176. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3177. auto maybe_layer = maybeCurrentDynamicLayer();
  3178. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3179. int64_t cur_level = maybe_layer->layerId();
  3180. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  3181. return at::_ops::conv_transpose1d::call(input, weight, bias, stride, padding, output_padding, groups, dilation);
  3182. }
  3183. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  3184. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  3185. std::optional<Tensor> bias_value;
  3186. std::optional<int64_t> bias_bdim;
  3187. if (bias) {
  3188. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  3189. }
  3190. auto results = batch_rule(input_value, input_bdim, weight_value, weight_bdim, bias_value, bias_bdim, stride, padding, output_padding, groups, dilation);
  3191. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  3192. }
  3193. template <typename batch_rule_t, batch_rule_t batch_rule>
  3194. at::Tensor conv_transpose2d_input_generated_plumbing(const at::Tensor & input, const at::Tensor & weight, const ::std::optional<at::Tensor> & bias, c10::SymIntArrayRef stride, c10::SymIntArrayRef padding, c10::SymIntArrayRef output_padding, c10::SymInt groups, c10::SymIntArrayRef dilation) {
  3195. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3196. auto maybe_layer = maybeCurrentDynamicLayer();
  3197. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3198. int64_t cur_level = maybe_layer->layerId();
  3199. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  3200. return at::_ops::conv_transpose2d_input::call(input, weight, bias, stride, padding, output_padding, groups, dilation);
  3201. }
  3202. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  3203. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  3204. std::optional<Tensor> bias_value;
  3205. std::optional<int64_t> bias_bdim;
  3206. if (bias) {
  3207. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  3208. }
  3209. auto results = batch_rule(input_value, input_bdim, weight_value, weight_bdim, bias_value, bias_bdim, stride, padding, output_padding, groups, dilation);
  3210. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  3211. }
  3212. template <typename batch_rule_t, batch_rule_t batch_rule>
  3213. at::Tensor conv_transpose3d_input_generated_plumbing(const at::Tensor & input, const at::Tensor & weight, const ::std::optional<at::Tensor> & bias, c10::SymIntArrayRef stride, c10::SymIntArrayRef padding, c10::SymIntArrayRef output_padding, c10::SymInt groups, c10::SymIntArrayRef dilation) {
  3214. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3215. auto maybe_layer = maybeCurrentDynamicLayer();
  3216. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3217. int64_t cur_level = maybe_layer->layerId();
  3218. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  3219. return at::_ops::conv_transpose3d_input::call(input, weight, bias, stride, padding, output_padding, groups, dilation);
  3220. }
  3221. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  3222. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  3223. std::optional<Tensor> bias_value;
  3224. std::optional<int64_t> bias_bdim;
  3225. if (bias) {
  3226. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  3227. }
  3228. auto results = batch_rule(input_value, input_bdim, weight_value, weight_bdim, bias_value, bias_bdim, stride, padding, output_padding, groups, dilation);
  3229. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  3230. }
  3231. template <typename batch_rule_t, batch_rule_t batch_rule>
  3232. at::Tensor copy_generated_plumbing(const at::Tensor & self, const at::Tensor & src, bool non_blocking) {
  3233. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3234. auto maybe_layer = maybeCurrentDynamicLayer();
  3235. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3236. int64_t cur_level = maybe_layer->layerId();
  3237. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(src, cur_level)) {
  3238. return at::_ops::copy::call(self, src, non_blocking);
  3239. }
  3240. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  3241. auto [src_value, src_bdim] = unwrapTensorAtLevel(src, cur_level);
  3242. auto results = batch_rule(self_value, self_bdim, src_value, src_bdim, non_blocking);
  3243. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  3244. }
  3245. template <typename batch_rule_t, batch_rule_t batch_rule>
  3246. at::Tensor & copy__generated_plumbing(at::Tensor & self, const at::Tensor & src, bool non_blocking) {
  3247. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3248. auto maybe_layer = maybeCurrentDynamicLayer();
  3249. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  3250. int64_t cur_level = maybe_layer->layerId();
  3251. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(src, cur_level)) {
  3252. return at::_ops::copy_::call(self, src, non_blocking);
  3253. }
  3254. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  3255. auto [src_value, src_bdim] = unwrapTensorAtLevel(src, cur_level);
  3256. batch_rule(self_value, self_bdim, src_value, src_bdim, non_blocking);
  3257. return self;
  3258. }
  3259. template <typename batch_rule_t, batch_rule_t batch_rule>
  3260. at::Tensor _copy_from_generated_plumbing(const at::Tensor & self, const at::Tensor & dst, bool non_blocking) {
  3261. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3262. auto maybe_layer = maybeCurrentDynamicLayer();
  3263. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3264. int64_t cur_level = maybe_layer->layerId();
  3265. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(dst, cur_level)) {
  3266. return at::_ops::_copy_from::call(self, dst, non_blocking);
  3267. }
  3268. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  3269. auto [dst_value, dst_bdim] = unwrapTensorAtLevel(dst, cur_level);
  3270. auto results = batch_rule(self_value, self_bdim, dst_value, dst_bdim, non_blocking);
  3271. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  3272. }
  3273. template <typename batch_rule_t, batch_rule_t batch_rule>
  3274. at::Tensor _copy_from_and_resize_generated_plumbing(const at::Tensor & self, const at::Tensor & dst) {
  3275. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3276. auto maybe_layer = maybeCurrentDynamicLayer();
  3277. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3278. int64_t cur_level = maybe_layer->layerId();
  3279. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(dst, cur_level)) {
  3280. return at::_ops::_copy_from_and_resize::call(self, dst);
  3281. }
  3282. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  3283. auto [dst_value, dst_bdim] = unwrapTensorAtLevel(dst, cur_level);
  3284. auto results = batch_rule(self_value, self_bdim, dst_value, dst_bdim);
  3285. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  3286. }
  3287. template <typename batch_rule_t, batch_rule_t batch_rule>
  3288. at::Tensor cos_generated_plumbing(const at::Tensor & self) {
  3289. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3290. auto maybe_layer = maybeCurrentDynamicLayer();
  3291. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3292. int64_t cur_level = maybe_layer->layerId();
  3293. if (!isBatchedAtLevel(self, cur_level)) {
  3294. return at::_ops::cos::call(self);
  3295. }
  3296. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  3297. auto results = batch_rule(self_value, self_bdim);
  3298. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  3299. }
  3300. template <typename batch_rule_t, batch_rule_t batch_rule>
  3301. at::Tensor & cos__generated_plumbing(at::Tensor & self) {
  3302. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3303. auto maybe_layer = maybeCurrentDynamicLayer();
  3304. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  3305. int64_t cur_level = maybe_layer->layerId();
  3306. if (!isBatchedAtLevel(self, cur_level)) {
  3307. return at::_ops::cos_::call(self);
  3308. }
  3309. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  3310. batch_rule(self_value, self_bdim);
  3311. return self;
  3312. }
  3313. template <typename batch_rule_t, batch_rule_t batch_rule>
  3314. at::Tensor cosh_generated_plumbing(const at::Tensor & self) {
  3315. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3316. auto maybe_layer = maybeCurrentDynamicLayer();
  3317. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3318. int64_t cur_level = maybe_layer->layerId();
  3319. if (!isBatchedAtLevel(self, cur_level)) {
  3320. return at::_ops::cosh::call(self);
  3321. }
  3322. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  3323. auto results = batch_rule(self_value, self_bdim);
  3324. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  3325. }
  3326. template <typename batch_rule_t, batch_rule_t batch_rule>
  3327. at::Tensor & cosh__generated_plumbing(at::Tensor & self) {
  3328. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3329. auto maybe_layer = maybeCurrentDynamicLayer();
  3330. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  3331. int64_t cur_level = maybe_layer->layerId();
  3332. if (!isBatchedAtLevel(self, cur_level)) {
  3333. return at::_ops::cosh_::call(self);
  3334. }
  3335. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  3336. batch_rule(self_value, self_bdim);
  3337. return self;
  3338. }
  3339. template <typename batch_rule_t, batch_rule_t batch_rule>
  3340. at::Tensor cosine_embedding_loss_generated_plumbing(const at::Tensor & input1, const at::Tensor & input2, const at::Tensor & target, double margin, int64_t reduction) {
  3341. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3342. auto maybe_layer = maybeCurrentDynamicLayer();
  3343. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3344. int64_t cur_level = maybe_layer->layerId();
  3345. if (!isBatchedAtLevel(input1, cur_level) && !isBatchedAtLevel(input2, cur_level) && !isBatchedAtLevel(target, cur_level)) {
  3346. return at::_ops::cosine_embedding_loss::call(input1, input2, target, margin, reduction);
  3347. }
  3348. auto [input1_value, input1_bdim] = unwrapTensorAtLevel(input1, cur_level);
  3349. auto [input2_value, input2_bdim] = unwrapTensorAtLevel(input2, cur_level);
  3350. auto [target_value, target_bdim] = unwrapTensorAtLevel(target, cur_level);
  3351. auto results = batch_rule(input1_value, input1_bdim, input2_value, input2_bdim, target_value, target_bdim, margin, reduction);
  3352. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  3353. }
  3354. template <typename batch_rule_t, batch_rule_t batch_rule>
  3355. at::Tensor count_nonzero_dim_IntList_generated_plumbing(const at::Tensor & self, at::IntArrayRef dim) {
  3356. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3357. auto maybe_layer = maybeCurrentDynamicLayer();
  3358. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3359. int64_t cur_level = maybe_layer->layerId();
  3360. if (!isBatchedAtLevel(self, cur_level)) {
  3361. return at::_ops::count_nonzero_dim_IntList::call(self, dim);
  3362. }
  3363. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  3364. auto results = batch_rule(self_value, self_bdim, dim);
  3365. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  3366. }
  3367. template <typename batch_rule_t, batch_rule_t batch_rule>
  3368. at::Tensor count_nonzero_generated_plumbing(const at::Tensor & self, ::std::optional<int64_t> dim) {
  3369. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3370. auto maybe_layer = maybeCurrentDynamicLayer();
  3371. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3372. int64_t cur_level = maybe_layer->layerId();
  3373. if (!isBatchedAtLevel(self, cur_level)) {
  3374. return at::_ops::count_nonzero::call(self, dim);
  3375. }
  3376. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  3377. auto results = batch_rule(self_value, self_bdim, dim);
  3378. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  3379. }
  3380. template <typename batch_rule_t, batch_rule_t batch_rule>
  3381. at::Tensor cov_generated_plumbing(const at::Tensor & self, int64_t correction, const ::std::optional<at::Tensor> & fweights, const ::std::optional<at::Tensor> & aweights) {
  3382. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3383. auto maybe_layer = maybeCurrentDynamicLayer();
  3384. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3385. int64_t cur_level = maybe_layer->layerId();
  3386. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(fweights, cur_level) && !isBatchedAtLevel(aweights, cur_level)) {
  3387. return at::_ops::cov::call(self, correction, fweights, aweights);
  3388. }
  3389. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  3390. std::optional<Tensor> fweights_value;
  3391. std::optional<int64_t> fweights_bdim;
  3392. if (fweights) {
  3393. std::tie(fweights_value, fweights_bdim) = unwrapTensorAtLevel(fweights.value(), cur_level);
  3394. }
  3395. std::optional<Tensor> aweights_value;
  3396. std::optional<int64_t> aweights_bdim;
  3397. if (aweights) {
  3398. std::tie(aweights_value, aweights_bdim) = unwrapTensorAtLevel(aweights.value(), cur_level);
  3399. }
  3400. auto results = batch_rule(self_value, self_bdim, correction, fweights_value, fweights_bdim, aweights_value, aweights_bdim);
  3401. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  3402. }
  3403. template <typename batch_rule_t, batch_rule_t batch_rule>
  3404. at::Tensor corrcoef_generated_plumbing(const at::Tensor & self) {
  3405. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3406. auto maybe_layer = maybeCurrentDynamicLayer();
  3407. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3408. int64_t cur_level = maybe_layer->layerId();
  3409. if (!isBatchedAtLevel(self, cur_level)) {
  3410. return at::_ops::corrcoef::call(self);
  3411. }
  3412. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  3413. auto results = batch_rule(self_value, self_bdim);
  3414. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  3415. }
  3416. template <typename batch_rule_t, batch_rule_t batch_rule>
  3417. at::Tensor cudnn_affine_grid_generator_generated_plumbing(const at::Tensor & theta, int64_t N, int64_t C, int64_t H, int64_t W) {
  3418. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3419. auto maybe_layer = maybeCurrentDynamicLayer();
  3420. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3421. int64_t cur_level = maybe_layer->layerId();
  3422. if (!isBatchedAtLevel(theta, cur_level)) {
  3423. return at::_ops::cudnn_affine_grid_generator::call(theta, N, C, H, W);
  3424. }
  3425. auto [theta_value, theta_bdim] = unwrapTensorAtLevel(theta, cur_level);
  3426. auto results = batch_rule(theta_value, theta_bdim, N, C, H, W);
  3427. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  3428. }
  3429. template <typename batch_rule_t, batch_rule_t batch_rule>
  3430. at::Tensor cudnn_affine_grid_generator_backward_generated_plumbing(const at::Tensor & grad, int64_t N, int64_t C, int64_t H, int64_t W) {
  3431. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3432. auto maybe_layer = maybeCurrentDynamicLayer();
  3433. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3434. int64_t cur_level = maybe_layer->layerId();
  3435. if (!isBatchedAtLevel(grad, cur_level)) {
  3436. return at::_ops::cudnn_affine_grid_generator_backward::call(grad, N, C, H, W);
  3437. }
  3438. auto [grad_value, grad_bdim] = unwrapTensorAtLevel(grad, cur_level);
  3439. auto results = batch_rule(grad_value, grad_bdim, N, C, H, W);
  3440. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  3441. }
  3442. template <typename batch_rule_t, batch_rule_t batch_rule>
  3443. ::std::tuple<at::Tensor,at::Tensor,at::Tensor,at::Tensor> cudnn_batch_norm_generated_plumbing(const at::Tensor & input, const at::Tensor & weight, const ::std::optional<at::Tensor> & bias, const ::std::optional<at::Tensor> & running_mean, const ::std::optional<at::Tensor> & running_var, bool training, double exponential_average_factor, double epsilon) {
  3444. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3445. auto maybe_layer = maybeCurrentDynamicLayer();
  3446. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3447. int64_t cur_level = maybe_layer->layerId();
  3448. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level) && !isBatchedAtLevel(running_mean, cur_level) && !isBatchedAtLevel(running_var, cur_level)) {
  3449. return at::_ops::cudnn_batch_norm::call(input, weight, bias, running_mean, running_var, training, exponential_average_factor, epsilon);
  3450. }
  3451. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  3452. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  3453. std::optional<Tensor> bias_value;
  3454. std::optional<int64_t> bias_bdim;
  3455. if (bias) {
  3456. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  3457. }
  3458. std::optional<Tensor> running_mean_value;
  3459. std::optional<int64_t> running_mean_bdim;
  3460. if (running_mean) {
  3461. std::tie(running_mean_value, running_mean_bdim) = unwrapTensorAtLevel(running_mean.value(), cur_level);
  3462. }
  3463. std::optional<Tensor> running_var_value;
  3464. std::optional<int64_t> running_var_bdim;
  3465. if (running_var) {
  3466. std::tie(running_var_value, running_var_bdim) = unwrapTensorAtLevel(running_var.value(), cur_level);
  3467. }
  3468. auto results = batch_rule(input_value, input_bdim, weight_value, weight_bdim, bias_value, bias_bdim, running_mean_value, running_mean_bdim, running_var_value, running_var_bdim, training, exponential_average_factor, epsilon);
  3469. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level), makeBatched(std::get<6>(results), std::get<7>(results), cur_level));
  3470. }
  3471. template <typename batch_rule_t, batch_rule_t batch_rule>
  3472. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> cudnn_batch_norm_backward_generated_plumbing(const at::Tensor & input, const at::Tensor & grad_output, const at::Tensor & weight, const ::std::optional<at::Tensor> & running_mean, const ::std::optional<at::Tensor> & running_var, const ::std::optional<at::Tensor> & save_mean, const ::std::optional<at::Tensor> & save_var, double epsilon, const at::Tensor & reserveSpace) {
  3473. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3474. auto maybe_layer = maybeCurrentDynamicLayer();
  3475. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3476. int64_t cur_level = maybe_layer->layerId();
  3477. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(running_mean, cur_level) && !isBatchedAtLevel(running_var, cur_level) && !isBatchedAtLevel(save_mean, cur_level) && !isBatchedAtLevel(save_var, cur_level) && !isBatchedAtLevel(reserveSpace, cur_level)) {
  3478. return at::_ops::cudnn_batch_norm_backward::call(input, grad_output, weight, running_mean, running_var, save_mean, save_var, epsilon, reserveSpace);
  3479. }
  3480. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  3481. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  3482. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  3483. auto [reserveSpace_value, reserveSpace_bdim] = unwrapTensorAtLevel(reserveSpace, cur_level);
  3484. std::optional<Tensor> running_mean_value;
  3485. std::optional<int64_t> running_mean_bdim;
  3486. if (running_mean) {
  3487. std::tie(running_mean_value, running_mean_bdim) = unwrapTensorAtLevel(running_mean.value(), cur_level);
  3488. }
  3489. std::optional<Tensor> running_var_value;
  3490. std::optional<int64_t> running_var_bdim;
  3491. if (running_var) {
  3492. std::tie(running_var_value, running_var_bdim) = unwrapTensorAtLevel(running_var.value(), cur_level);
  3493. }
  3494. std::optional<Tensor> save_mean_value;
  3495. std::optional<int64_t> save_mean_bdim;
  3496. if (save_mean) {
  3497. std::tie(save_mean_value, save_mean_bdim) = unwrapTensorAtLevel(save_mean.value(), cur_level);
  3498. }
  3499. std::optional<Tensor> save_var_value;
  3500. std::optional<int64_t> save_var_bdim;
  3501. if (save_var) {
  3502. std::tie(save_var_value, save_var_bdim) = unwrapTensorAtLevel(save_var.value(), cur_level);
  3503. }
  3504. auto results = batch_rule(input_value, input_bdim, grad_output_value, grad_output_bdim, weight_value, weight_bdim, running_mean_value, running_mean_bdim, running_var_value, running_var_bdim, save_mean_value, save_mean_bdim, save_var_value, save_var_bdim, epsilon, reserveSpace_value, reserveSpace_bdim);
  3505. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  3506. }
  3507. template <typename batch_rule_t, batch_rule_t batch_rule>
  3508. at::Tensor cudnn_convolution_generated_plumbing(const at::Tensor & self, const at::Tensor & weight, c10::SymIntArrayRef padding, c10::SymIntArrayRef stride, c10::SymIntArrayRef dilation, c10::SymInt groups, bool benchmark, bool deterministic, bool allow_tf32) {
  3509. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3510. auto maybe_layer = maybeCurrentDynamicLayer();
  3511. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3512. int64_t cur_level = maybe_layer->layerId();
  3513. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(weight, cur_level)) {
  3514. return at::_ops::cudnn_convolution::call(self, weight, padding, stride, dilation, groups, benchmark, deterministic, allow_tf32);
  3515. }
  3516. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  3517. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  3518. auto results = batch_rule(self_value, self_bdim, weight_value, weight_bdim, padding, stride, dilation, groups, benchmark, deterministic, allow_tf32);
  3519. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  3520. }
  3521. template <typename batch_rule_t, batch_rule_t batch_rule>
  3522. at::Tensor cudnn_convolution_transpose_generated_plumbing(const at::Tensor & self, const at::Tensor & weight, c10::SymIntArrayRef padding, c10::SymIntArrayRef output_padding, c10::SymIntArrayRef stride, c10::SymIntArrayRef dilation, c10::SymInt groups, bool benchmark, bool deterministic, bool allow_tf32) {
  3523. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3524. auto maybe_layer = maybeCurrentDynamicLayer();
  3525. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3526. int64_t cur_level = maybe_layer->layerId();
  3527. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(weight, cur_level)) {
  3528. return at::_ops::cudnn_convolution_transpose::call(self, weight, padding, output_padding, stride, dilation, groups, benchmark, deterministic, allow_tf32);
  3529. }
  3530. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  3531. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  3532. auto results = batch_rule(self_value, self_bdim, weight_value, weight_bdim, padding, output_padding, stride, dilation, groups, benchmark, deterministic, allow_tf32);
  3533. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  3534. }
  3535. template <typename batch_rule_t, batch_rule_t batch_rule>
  3536. at::Tensor _mps_convolution_transpose_generated_plumbing(const at::Tensor & self, const at::Tensor & weight, c10::SymIntArrayRef padding, c10::SymIntArrayRef output_padding, c10::SymIntArrayRef stride, c10::SymIntArrayRef dilation, c10::SymInt groups) {
  3537. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3538. auto maybe_layer = maybeCurrentDynamicLayer();
  3539. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3540. int64_t cur_level = maybe_layer->layerId();
  3541. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(weight, cur_level)) {
  3542. return at::_ops::_mps_convolution_transpose::call(self, weight, padding, output_padding, stride, dilation, groups);
  3543. }
  3544. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  3545. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  3546. auto results = batch_rule(self_value, self_bdim, weight_value, weight_bdim, padding, output_padding, stride, dilation, groups);
  3547. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  3548. }
  3549. template <typename batch_rule_t, batch_rule_t batch_rule>
  3550. ::std::tuple<at::Tensor,at::Tensor> mps_convolution_transpose_backward_generated_plumbing(const at::Tensor & self, const at::Tensor & grad_output, const at::Tensor & weight, c10::SymIntArrayRef padding, c10::SymIntArrayRef output_padding, c10::SymIntArrayRef stride, c10::SymIntArrayRef dilation, c10::SymInt groups, ::std::array<bool,2> output_mask) {
  3551. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3552. auto maybe_layer = maybeCurrentDynamicLayer();
  3553. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3554. int64_t cur_level = maybe_layer->layerId();
  3555. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(weight, cur_level)) {
  3556. return at::_ops::mps_convolution_transpose_backward::call(self, grad_output, weight, padding, output_padding, stride, dilation, groups, output_mask);
  3557. }
  3558. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  3559. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  3560. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  3561. auto results = batch_rule(self_value, self_bdim, grad_output_value, grad_output_bdim, weight_value, weight_bdim, padding, output_padding, stride, dilation, groups, output_mask);
  3562. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  3563. }
  3564. template <typename batch_rule_t, batch_rule_t batch_rule>
  3565. at::Tensor cudnn_convolution_relu_generated_plumbing(const at::Tensor & self, const at::Tensor & weight, const ::std::optional<at::Tensor> & bias, c10::SymIntArrayRef stride, c10::SymIntArrayRef padding, c10::SymIntArrayRef dilation, c10::SymInt groups) {
  3566. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3567. auto maybe_layer = maybeCurrentDynamicLayer();
  3568. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3569. int64_t cur_level = maybe_layer->layerId();
  3570. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  3571. return at::_ops::cudnn_convolution_relu::call(self, weight, bias, stride, padding, dilation, groups);
  3572. }
  3573. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  3574. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  3575. std::optional<Tensor> bias_value;
  3576. std::optional<int64_t> bias_bdim;
  3577. if (bias) {
  3578. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  3579. }
  3580. auto results = batch_rule(self_value, self_bdim, weight_value, weight_bdim, bias_value, bias_bdim, stride, padding, dilation, groups);
  3581. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  3582. }
  3583. template <typename batch_rule_t, batch_rule_t batch_rule>
  3584. at::Tensor cudnn_convolution_add_relu_generated_plumbing(const at::Tensor & self, const at::Tensor & weight, const at::Tensor & z, const ::std::optional<at::Scalar> & alpha, const ::std::optional<at::Tensor> & bias, c10::SymIntArrayRef stride, c10::SymIntArrayRef padding, c10::SymIntArrayRef dilation, c10::SymInt groups) {
  3585. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3586. auto maybe_layer = maybeCurrentDynamicLayer();
  3587. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3588. int64_t cur_level = maybe_layer->layerId();
  3589. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(z, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  3590. return at::_ops::cudnn_convolution_add_relu::call(self, weight, z, alpha, bias, stride, padding, dilation, groups);
  3591. }
  3592. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  3593. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  3594. auto [z_value, z_bdim] = unwrapTensorAtLevel(z, cur_level);
  3595. std::optional<Tensor> bias_value;
  3596. std::optional<int64_t> bias_bdim;
  3597. if (bias) {
  3598. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  3599. }
  3600. auto results = batch_rule(self_value, self_bdim, weight_value, weight_bdim, z_value, z_bdim, alpha, bias_value, bias_bdim, stride, padding, dilation, groups);
  3601. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  3602. }
  3603. template <typename batch_rule_t, batch_rule_t batch_rule>
  3604. at::Tensor cudnn_grid_sampler_generated_plumbing(const at::Tensor & self, const at::Tensor & grid) {
  3605. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3606. auto maybe_layer = maybeCurrentDynamicLayer();
  3607. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3608. int64_t cur_level = maybe_layer->layerId();
  3609. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(grid, cur_level)) {
  3610. return at::_ops::cudnn_grid_sampler::call(self, grid);
  3611. }
  3612. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  3613. auto [grid_value, grid_bdim] = unwrapTensorAtLevel(grid, cur_level);
  3614. auto results = batch_rule(self_value, self_bdim, grid_value, grid_bdim);
  3615. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  3616. }
  3617. template <typename batch_rule_t, batch_rule_t batch_rule>
  3618. ::std::tuple<at::Tensor,at::Tensor> cudnn_grid_sampler_backward_generated_plumbing(const at::Tensor & self, const at::Tensor & grid, const at::Tensor & grad_output) {
  3619. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3620. auto maybe_layer = maybeCurrentDynamicLayer();
  3621. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3622. int64_t cur_level = maybe_layer->layerId();
  3623. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(grid, cur_level) && !isBatchedAtLevel(grad_output, cur_level)) {
  3624. return at::_ops::cudnn_grid_sampler_backward::call(self, grid, grad_output);
  3625. }
  3626. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  3627. auto [grid_value, grid_bdim] = unwrapTensorAtLevel(grid, cur_level);
  3628. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  3629. auto results = batch_rule(self_value, self_bdim, grid_value, grid_bdim, grad_output_value, grad_output_bdim);
  3630. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  3631. }
  3632. template <typename batch_rule_t, batch_rule_t batch_rule>
  3633. ::std::tuple<at::Tensor,at::Tensor> cummax_generated_plumbing(const at::Tensor & self, int64_t dim) {
  3634. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3635. auto maybe_layer = maybeCurrentDynamicLayer();
  3636. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3637. int64_t cur_level = maybe_layer->layerId();
  3638. if (!isBatchedAtLevel(self, cur_level)) {
  3639. return at::_ops::cummax::call(self, dim);
  3640. }
  3641. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  3642. auto results = batch_rule(self_value, self_bdim, dim);
  3643. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  3644. }
  3645. template <typename batch_rule_t, batch_rule_t batch_rule>
  3646. ::std::tuple<at::Tensor,at::Tensor> cummax_dimname_generated_plumbing(const at::Tensor & self, at::Dimname dim) {
  3647. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3648. auto maybe_layer = maybeCurrentDynamicLayer();
  3649. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3650. int64_t cur_level = maybe_layer->layerId();
  3651. if (!isBatchedAtLevel(self, cur_level)) {
  3652. return at::_ops::cummax_dimname::call(self, dim);
  3653. }
  3654. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  3655. auto results = batch_rule(self_value, self_bdim, dim);
  3656. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  3657. }
  3658. template <typename batch_rule_t, batch_rule_t batch_rule>
  3659. void _cummax_helper_generated_plumbing(const at::Tensor & self, at::Tensor & values, at::Tensor & indices, int64_t dim) {
  3660. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3661. auto maybe_layer = maybeCurrentDynamicLayer();
  3662. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  3663. int64_t cur_level = maybe_layer->layerId();
  3664. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(values, cur_level) && !isBatchedAtLevel(indices, cur_level)) {
  3665. return at::_ops::_cummax_helper::call(self, values, indices, dim);
  3666. }
  3667. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  3668. auto [values_value, values_bdim] = unwrapTensorAtLevel(values, cur_level);
  3669. auto [indices_value, indices_bdim] = unwrapTensorAtLevel(indices, cur_level);
  3670. batch_rule(self_value, self_bdim, values_value, values_bdim, indices_value, indices_bdim, dim);
  3671. }
  3672. template <typename batch_rule_t, batch_rule_t batch_rule>
  3673. ::std::tuple<at::Tensor,at::Tensor> cummin_generated_plumbing(const at::Tensor & self, int64_t dim) {
  3674. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3675. auto maybe_layer = maybeCurrentDynamicLayer();
  3676. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3677. int64_t cur_level = maybe_layer->layerId();
  3678. if (!isBatchedAtLevel(self, cur_level)) {
  3679. return at::_ops::cummin::call(self, dim);
  3680. }
  3681. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  3682. auto results = batch_rule(self_value, self_bdim, dim);
  3683. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  3684. }
  3685. template <typename batch_rule_t, batch_rule_t batch_rule>
  3686. ::std::tuple<at::Tensor,at::Tensor> cummin_dimname_generated_plumbing(const at::Tensor & self, at::Dimname dim) {
  3687. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3688. auto maybe_layer = maybeCurrentDynamicLayer();
  3689. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3690. int64_t cur_level = maybe_layer->layerId();
  3691. if (!isBatchedAtLevel(self, cur_level)) {
  3692. return at::_ops::cummin_dimname::call(self, dim);
  3693. }
  3694. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  3695. auto results = batch_rule(self_value, self_bdim, dim);
  3696. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  3697. }
  3698. template <typename batch_rule_t, batch_rule_t batch_rule>
  3699. void _cummin_helper_generated_plumbing(const at::Tensor & self, at::Tensor & values, at::Tensor & indices, int64_t dim) {
  3700. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3701. auto maybe_layer = maybeCurrentDynamicLayer();
  3702. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  3703. int64_t cur_level = maybe_layer->layerId();
  3704. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(values, cur_level) && !isBatchedAtLevel(indices, cur_level)) {
  3705. return at::_ops::_cummin_helper::call(self, values, indices, dim);
  3706. }
  3707. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  3708. auto [values_value, values_bdim] = unwrapTensorAtLevel(values, cur_level);
  3709. auto [indices_value, indices_bdim] = unwrapTensorAtLevel(indices, cur_level);
  3710. batch_rule(self_value, self_bdim, values_value, values_bdim, indices_value, indices_bdim, dim);
  3711. }
  3712. template <typename batch_rule_t, batch_rule_t batch_rule>
  3713. at::Tensor cummaxmin_backward_generated_plumbing(const at::Tensor & grad, const at::Tensor & input, const at::Tensor & indices, int64_t dim) {
  3714. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3715. auto maybe_layer = maybeCurrentDynamicLayer();
  3716. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3717. int64_t cur_level = maybe_layer->layerId();
  3718. if (!isBatchedAtLevel(grad, cur_level) && !isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(indices, cur_level)) {
  3719. return at::_ops::cummaxmin_backward::call(grad, input, indices, dim);
  3720. }
  3721. auto [grad_value, grad_bdim] = unwrapTensorAtLevel(grad, cur_level);
  3722. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  3723. auto [indices_value, indices_bdim] = unwrapTensorAtLevel(indices, cur_level);
  3724. auto results = batch_rule(grad_value, grad_bdim, input_value, input_bdim, indices_value, indices_bdim, dim);
  3725. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  3726. }
  3727. template <typename batch_rule_t, batch_rule_t batch_rule>
  3728. at::Tensor cumprod_generated_plumbing(const at::Tensor & self, int64_t dim, ::std::optional<at::ScalarType> dtype) {
  3729. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3730. auto maybe_layer = maybeCurrentDynamicLayer();
  3731. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3732. int64_t cur_level = maybe_layer->layerId();
  3733. if (!isBatchedAtLevel(self, cur_level)) {
  3734. return at::_ops::cumprod::call(self, dim, dtype);
  3735. }
  3736. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  3737. auto results = batch_rule(self_value, self_bdim, dim, dtype);
  3738. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  3739. }
  3740. template <typename batch_rule_t, batch_rule_t batch_rule>
  3741. at::Tensor & cumprod__generated_plumbing(at::Tensor & self, int64_t dim, ::std::optional<at::ScalarType> dtype) {
  3742. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3743. auto maybe_layer = maybeCurrentDynamicLayer();
  3744. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  3745. int64_t cur_level = maybe_layer->layerId();
  3746. if (!isBatchedAtLevel(self, cur_level)) {
  3747. return at::_ops::cumprod_::call(self, dim, dtype);
  3748. }
  3749. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  3750. batch_rule(self_value, self_bdim, dim, dtype);
  3751. return self;
  3752. }
  3753. template <typename batch_rule_t, batch_rule_t batch_rule>
  3754. at::Tensor cumprod_dimname_generated_plumbing(const at::Tensor & self, at::Dimname dim, ::std::optional<at::ScalarType> dtype) {
  3755. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3756. auto maybe_layer = maybeCurrentDynamicLayer();
  3757. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3758. int64_t cur_level = maybe_layer->layerId();
  3759. if (!isBatchedAtLevel(self, cur_level)) {
  3760. return at::_ops::cumprod_dimname::call(self, dim, dtype);
  3761. }
  3762. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  3763. auto results = batch_rule(self_value, self_bdim, dim, dtype);
  3764. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  3765. }
  3766. template <typename batch_rule_t, batch_rule_t batch_rule>
  3767. at::Tensor & cumprod__dimname_generated_plumbing(at::Tensor & self, at::Dimname dim, ::std::optional<at::ScalarType> dtype) {
  3768. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3769. auto maybe_layer = maybeCurrentDynamicLayer();
  3770. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  3771. int64_t cur_level = maybe_layer->layerId();
  3772. if (!isBatchedAtLevel(self, cur_level)) {
  3773. return at::_ops::cumprod__dimname::call(self, dim, dtype);
  3774. }
  3775. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  3776. batch_rule(self_value, self_bdim, dim, dtype);
  3777. return self;
  3778. }
  3779. template <typename batch_rule_t, batch_rule_t batch_rule>
  3780. at::Tensor cumprod_backward_generated_plumbing(const at::Tensor & grad, const at::Tensor & input, int64_t dim, const at::Tensor & output) {
  3781. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3782. auto maybe_layer = maybeCurrentDynamicLayer();
  3783. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3784. int64_t cur_level = maybe_layer->layerId();
  3785. if (!isBatchedAtLevel(grad, cur_level) && !isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(output, cur_level)) {
  3786. return at::_ops::cumprod_backward::call(grad, input, dim, output);
  3787. }
  3788. auto [grad_value, grad_bdim] = unwrapTensorAtLevel(grad, cur_level);
  3789. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  3790. auto [output_value, output_bdim] = unwrapTensorAtLevel(output, cur_level);
  3791. auto results = batch_rule(grad_value, grad_bdim, input_value, input_bdim, dim, output_value, output_bdim);
  3792. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  3793. }
  3794. template <typename batch_rule_t, batch_rule_t batch_rule>
  3795. at::Tensor cumsum_generated_plumbing(const at::Tensor & self, int64_t dim, ::std::optional<at::ScalarType> dtype) {
  3796. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3797. auto maybe_layer = maybeCurrentDynamicLayer();
  3798. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3799. int64_t cur_level = maybe_layer->layerId();
  3800. if (!isBatchedAtLevel(self, cur_level)) {
  3801. return at::_ops::cumsum::call(self, dim, dtype);
  3802. }
  3803. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  3804. auto results = batch_rule(self_value, self_bdim, dim, dtype);
  3805. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  3806. }
  3807. template <typename batch_rule_t, batch_rule_t batch_rule>
  3808. at::Tensor & cumsum__generated_plumbing(at::Tensor & self, int64_t dim, ::std::optional<at::ScalarType> dtype) {
  3809. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3810. auto maybe_layer = maybeCurrentDynamicLayer();
  3811. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  3812. int64_t cur_level = maybe_layer->layerId();
  3813. if (!isBatchedAtLevel(self, cur_level)) {
  3814. return at::_ops::cumsum_::call(self, dim, dtype);
  3815. }
  3816. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  3817. batch_rule(self_value, self_bdim, dim, dtype);
  3818. return self;
  3819. }
  3820. template <typename batch_rule_t, batch_rule_t batch_rule>
  3821. at::Tensor cumsum_dimname_generated_plumbing(const at::Tensor & self, at::Dimname dim, ::std::optional<at::ScalarType> dtype) {
  3822. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3823. auto maybe_layer = maybeCurrentDynamicLayer();
  3824. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3825. int64_t cur_level = maybe_layer->layerId();
  3826. if (!isBatchedAtLevel(self, cur_level)) {
  3827. return at::_ops::cumsum_dimname::call(self, dim, dtype);
  3828. }
  3829. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  3830. auto results = batch_rule(self_value, self_bdim, dim, dtype);
  3831. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  3832. }
  3833. template <typename batch_rule_t, batch_rule_t batch_rule>
  3834. at::Tensor & cumsum__dimname_generated_plumbing(at::Tensor & self, at::Dimname dim, ::std::optional<at::ScalarType> dtype) {
  3835. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3836. auto maybe_layer = maybeCurrentDynamicLayer();
  3837. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  3838. int64_t cur_level = maybe_layer->layerId();
  3839. if (!isBatchedAtLevel(self, cur_level)) {
  3840. return at::_ops::cumsum__dimname::call(self, dim, dtype);
  3841. }
  3842. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  3843. batch_rule(self_value, self_bdim, dim, dtype);
  3844. return self;
  3845. }
  3846. template <typename batch_rule_t, batch_rule_t batch_rule>
  3847. at::Tensor cumulative_trapezoid_x_generated_plumbing(const at::Tensor & y, const at::Tensor & x, int64_t dim) {
  3848. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3849. auto maybe_layer = maybeCurrentDynamicLayer();
  3850. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3851. int64_t cur_level = maybe_layer->layerId();
  3852. if (!isBatchedAtLevel(y, cur_level) && !isBatchedAtLevel(x, cur_level)) {
  3853. return at::_ops::cumulative_trapezoid_x::call(y, x, dim);
  3854. }
  3855. auto [y_value, y_bdim] = unwrapTensorAtLevel(y, cur_level);
  3856. auto [x_value, x_bdim] = unwrapTensorAtLevel(x, cur_level);
  3857. auto results = batch_rule(y_value, y_bdim, x_value, x_bdim, dim);
  3858. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  3859. }
  3860. template <typename batch_rule_t, batch_rule_t batch_rule>
  3861. at::Tensor cumulative_trapezoid_dx_generated_plumbing(const at::Tensor & y, const at::Scalar & dx, int64_t dim) {
  3862. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3863. auto maybe_layer = maybeCurrentDynamicLayer();
  3864. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3865. int64_t cur_level = maybe_layer->layerId();
  3866. if (!isBatchedAtLevel(y, cur_level)) {
  3867. return at::_ops::cumulative_trapezoid_dx::call(y, dx, dim);
  3868. }
  3869. auto [y_value, y_bdim] = unwrapTensorAtLevel(y, cur_level);
  3870. auto results = batch_rule(y_value, y_bdim, dx, dim);
  3871. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  3872. }
  3873. template <typename batch_rule_t, batch_rule_t batch_rule>
  3874. at::Tensor ctc_loss_IntList_generated_plumbing(const at::Tensor & log_probs, const at::Tensor & targets, at::IntArrayRef input_lengths, at::IntArrayRef target_lengths, int64_t blank, int64_t reduction, bool zero_infinity) {
  3875. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3876. auto maybe_layer = maybeCurrentDynamicLayer();
  3877. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3878. int64_t cur_level = maybe_layer->layerId();
  3879. if (!isBatchedAtLevel(log_probs, cur_level) && !isBatchedAtLevel(targets, cur_level)) {
  3880. return at::_ops::ctc_loss_IntList::call(log_probs, targets, input_lengths, target_lengths, blank, reduction, zero_infinity);
  3881. }
  3882. auto [log_probs_value, log_probs_bdim] = unwrapTensorAtLevel(log_probs, cur_level);
  3883. auto [targets_value, targets_bdim] = unwrapTensorAtLevel(targets, cur_level);
  3884. auto results = batch_rule(log_probs_value, log_probs_bdim, targets_value, targets_bdim, input_lengths, target_lengths, blank, reduction, zero_infinity);
  3885. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  3886. }
  3887. template <typename batch_rule_t, batch_rule_t batch_rule>
  3888. at::Tensor ctc_loss_Tensor_generated_plumbing(const at::Tensor & log_probs, const at::Tensor & targets, const at::Tensor & input_lengths, const at::Tensor & target_lengths, int64_t blank, int64_t reduction, bool zero_infinity) {
  3889. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3890. auto maybe_layer = maybeCurrentDynamicLayer();
  3891. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3892. int64_t cur_level = maybe_layer->layerId();
  3893. if (!isBatchedAtLevel(log_probs, cur_level) && !isBatchedAtLevel(targets, cur_level) && !isBatchedAtLevel(input_lengths, cur_level) && !isBatchedAtLevel(target_lengths, cur_level)) {
  3894. return at::_ops::ctc_loss_Tensor::call(log_probs, targets, input_lengths, target_lengths, blank, reduction, zero_infinity);
  3895. }
  3896. auto [log_probs_value, log_probs_bdim] = unwrapTensorAtLevel(log_probs, cur_level);
  3897. auto [targets_value, targets_bdim] = unwrapTensorAtLevel(targets, cur_level);
  3898. auto [input_lengths_value, input_lengths_bdim] = unwrapTensorAtLevel(input_lengths, cur_level);
  3899. auto [target_lengths_value, target_lengths_bdim] = unwrapTensorAtLevel(target_lengths, cur_level);
  3900. auto results = batch_rule(log_probs_value, log_probs_bdim, targets_value, targets_bdim, input_lengths_value, input_lengths_bdim, target_lengths_value, target_lengths_bdim, blank, reduction, zero_infinity);
  3901. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  3902. }
  3903. template <typename batch_rule_t, batch_rule_t batch_rule>
  3904. ::std::tuple<at::Tensor,at::Tensor> _ctc_loss_generated_plumbing(const at::Tensor & log_probs, const at::Tensor & targets, at::IntArrayRef input_lengths, at::IntArrayRef target_lengths, int64_t blank, bool zero_infinity) {
  3905. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3906. auto maybe_layer = maybeCurrentDynamicLayer();
  3907. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3908. int64_t cur_level = maybe_layer->layerId();
  3909. if (!isBatchedAtLevel(log_probs, cur_level) && !isBatchedAtLevel(targets, cur_level)) {
  3910. return at::_ops::_ctc_loss::call(log_probs, targets, input_lengths, target_lengths, blank, zero_infinity);
  3911. }
  3912. auto [log_probs_value, log_probs_bdim] = unwrapTensorAtLevel(log_probs, cur_level);
  3913. auto [targets_value, targets_bdim] = unwrapTensorAtLevel(targets, cur_level);
  3914. auto results = batch_rule(log_probs_value, log_probs_bdim, targets_value, targets_bdim, input_lengths, target_lengths, blank, zero_infinity);
  3915. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  3916. }
  3917. template <typename batch_rule_t, batch_rule_t batch_rule>
  3918. ::std::tuple<at::Tensor,at::Tensor> _ctc_loss_Tensor_generated_plumbing(const at::Tensor & log_probs, const at::Tensor & targets, const at::Tensor & input_lengths, const at::Tensor & target_lengths, int64_t blank, bool zero_infinity) {
  3919. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3920. auto maybe_layer = maybeCurrentDynamicLayer();
  3921. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3922. int64_t cur_level = maybe_layer->layerId();
  3923. if (!isBatchedAtLevel(log_probs, cur_level) && !isBatchedAtLevel(targets, cur_level) && !isBatchedAtLevel(input_lengths, cur_level) && !isBatchedAtLevel(target_lengths, cur_level)) {
  3924. return at::_ops::_ctc_loss_Tensor::call(log_probs, targets, input_lengths, target_lengths, blank, zero_infinity);
  3925. }
  3926. auto [log_probs_value, log_probs_bdim] = unwrapTensorAtLevel(log_probs, cur_level);
  3927. auto [targets_value, targets_bdim] = unwrapTensorAtLevel(targets, cur_level);
  3928. auto [input_lengths_value, input_lengths_bdim] = unwrapTensorAtLevel(input_lengths, cur_level);
  3929. auto [target_lengths_value, target_lengths_bdim] = unwrapTensorAtLevel(target_lengths, cur_level);
  3930. auto results = batch_rule(log_probs_value, log_probs_bdim, targets_value, targets_bdim, input_lengths_value, input_lengths_bdim, target_lengths_value, target_lengths_bdim, blank, zero_infinity);
  3931. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  3932. }
  3933. template <typename batch_rule_t, batch_rule_t batch_rule>
  3934. at::Tensor _ctc_loss_backward_generated_plumbing(const at::Tensor & grad, const at::Tensor & log_probs, const at::Tensor & targets, at::IntArrayRef input_lengths, at::IntArrayRef target_lengths, const at::Tensor & neg_log_likelihood, const at::Tensor & log_alpha, int64_t blank, bool zero_infinity) {
  3935. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3936. auto maybe_layer = maybeCurrentDynamicLayer();
  3937. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3938. int64_t cur_level = maybe_layer->layerId();
  3939. if (!isBatchedAtLevel(grad, cur_level) && !isBatchedAtLevel(log_probs, cur_level) && !isBatchedAtLevel(targets, cur_level) && !isBatchedAtLevel(neg_log_likelihood, cur_level) && !isBatchedAtLevel(log_alpha, cur_level)) {
  3940. return at::_ops::_ctc_loss_backward::call(grad, log_probs, targets, input_lengths, target_lengths, neg_log_likelihood, log_alpha, blank, zero_infinity);
  3941. }
  3942. auto [grad_value, grad_bdim] = unwrapTensorAtLevel(grad, cur_level);
  3943. auto [log_probs_value, log_probs_bdim] = unwrapTensorAtLevel(log_probs, cur_level);
  3944. auto [targets_value, targets_bdim] = unwrapTensorAtLevel(targets, cur_level);
  3945. auto [neg_log_likelihood_value, neg_log_likelihood_bdim] = unwrapTensorAtLevel(neg_log_likelihood, cur_level);
  3946. auto [log_alpha_value, log_alpha_bdim] = unwrapTensorAtLevel(log_alpha, cur_level);
  3947. auto results = batch_rule(grad_value, grad_bdim, log_probs_value, log_probs_bdim, targets_value, targets_bdim, input_lengths, target_lengths, neg_log_likelihood_value, neg_log_likelihood_bdim, log_alpha_value, log_alpha_bdim, blank, zero_infinity);
  3948. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  3949. }
  3950. template <typename batch_rule_t, batch_rule_t batch_rule>
  3951. at::Tensor _ctc_loss_backward_Tensor_generated_plumbing(const at::Tensor & grad, const at::Tensor & log_probs, const at::Tensor & targets, const at::Tensor & input_lengths, const at::Tensor & target_lengths, const at::Tensor & neg_log_likelihood, const at::Tensor & log_alpha, int64_t blank, bool zero_infinity) {
  3952. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3953. auto maybe_layer = maybeCurrentDynamicLayer();
  3954. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3955. int64_t cur_level = maybe_layer->layerId();
  3956. if (!isBatchedAtLevel(grad, cur_level) && !isBatchedAtLevel(log_probs, cur_level) && !isBatchedAtLevel(targets, cur_level) && !isBatchedAtLevel(input_lengths, cur_level) && !isBatchedAtLevel(target_lengths, cur_level) && !isBatchedAtLevel(neg_log_likelihood, cur_level) && !isBatchedAtLevel(log_alpha, cur_level)) {
  3957. return at::_ops::_ctc_loss_backward_Tensor::call(grad, log_probs, targets, input_lengths, target_lengths, neg_log_likelihood, log_alpha, blank, zero_infinity);
  3958. }
  3959. auto [grad_value, grad_bdim] = unwrapTensorAtLevel(grad, cur_level);
  3960. auto [log_probs_value, log_probs_bdim] = unwrapTensorAtLevel(log_probs, cur_level);
  3961. auto [targets_value, targets_bdim] = unwrapTensorAtLevel(targets, cur_level);
  3962. auto [input_lengths_value, input_lengths_bdim] = unwrapTensorAtLevel(input_lengths, cur_level);
  3963. auto [target_lengths_value, target_lengths_bdim] = unwrapTensorAtLevel(target_lengths, cur_level);
  3964. auto [neg_log_likelihood_value, neg_log_likelihood_bdim] = unwrapTensorAtLevel(neg_log_likelihood, cur_level);
  3965. auto [log_alpha_value, log_alpha_bdim] = unwrapTensorAtLevel(log_alpha, cur_level);
  3966. auto results = batch_rule(grad_value, grad_bdim, log_probs_value, log_probs_bdim, targets_value, targets_bdim, input_lengths_value, input_lengths_bdim, target_lengths_value, target_lengths_bdim, neg_log_likelihood_value, neg_log_likelihood_bdim, log_alpha_value, log_alpha_bdim, blank, zero_infinity);
  3967. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  3968. }
  3969. template <typename batch_rule_t, batch_rule_t batch_rule>
  3970. at::Tensor diag_embed_generated_plumbing(const at::Tensor & self, int64_t offset, int64_t dim1, int64_t dim2) {
  3971. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3972. auto maybe_layer = maybeCurrentDynamicLayer();
  3973. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3974. int64_t cur_level = maybe_layer->layerId();
  3975. if (!isBatchedAtLevel(self, cur_level)) {
  3976. return at::_ops::diag_embed::call(self, offset, dim1, dim2);
  3977. }
  3978. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  3979. auto results = batch_rule(self_value, self_bdim, offset, dim1, dim2);
  3980. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  3981. }
  3982. template <typename batch_rule_t, batch_rule_t batch_rule>
  3983. at::Tensor diagflat_generated_plumbing(const at::Tensor & self, int64_t offset) {
  3984. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3985. auto maybe_layer = maybeCurrentDynamicLayer();
  3986. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  3987. int64_t cur_level = maybe_layer->layerId();
  3988. if (!isBatchedAtLevel(self, cur_level)) {
  3989. return at::_ops::diagflat::call(self, offset);
  3990. }
  3991. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  3992. auto results = batch_rule(self_value, self_bdim, offset);
  3993. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  3994. }
  3995. template <typename batch_rule_t, batch_rule_t batch_rule>
  3996. at::Tensor diagonal_generated_plumbing(const at::Tensor & self, int64_t offset, int64_t dim1, int64_t dim2) {
  3997. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  3998. auto maybe_layer = maybeCurrentDynamicLayer();
  3999. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4000. int64_t cur_level = maybe_layer->layerId();
  4001. if (!isBatchedAtLevel(self, cur_level)) {
  4002. return at::_ops::diagonal::call(self, offset, dim1, dim2);
  4003. }
  4004. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4005. auto results = batch_rule(self_value, self_bdim, offset, dim1, dim2);
  4006. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  4007. }
  4008. template <typename batch_rule_t, batch_rule_t batch_rule>
  4009. at::Tensor linalg_diagonal_generated_plumbing(const at::Tensor & A, int64_t offset, int64_t dim1, int64_t dim2) {
  4010. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4011. auto maybe_layer = maybeCurrentDynamicLayer();
  4012. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4013. int64_t cur_level = maybe_layer->layerId();
  4014. if (!isBatchedAtLevel(A, cur_level)) {
  4015. return at::_ops::linalg_diagonal::call(A, offset, dim1, dim2);
  4016. }
  4017. auto [A_value, A_bdim] = unwrapTensorAtLevel(A, cur_level);
  4018. auto results = batch_rule(A_value, A_bdim, offset, dim1, dim2);
  4019. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  4020. }
  4021. template <typename batch_rule_t, batch_rule_t batch_rule>
  4022. at::Tensor diagonal_Dimname_generated_plumbing(const at::Tensor & self, at::Dimname outdim, at::Dimname dim1, at::Dimname dim2, int64_t offset) {
  4023. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4024. auto maybe_layer = maybeCurrentDynamicLayer();
  4025. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4026. int64_t cur_level = maybe_layer->layerId();
  4027. if (!isBatchedAtLevel(self, cur_level)) {
  4028. return at::_ops::diagonal_Dimname::call(self, outdim, dim1, dim2, offset);
  4029. }
  4030. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4031. auto results = batch_rule(self_value, self_bdim, outdim, dim1, dim2, offset);
  4032. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  4033. }
  4034. template <typename batch_rule_t, batch_rule_t batch_rule>
  4035. at::Tensor diagonal_backward_generated_plumbing(const at::Tensor & grad_output, c10::SymIntArrayRef input_sizes, int64_t offset, int64_t dim1, int64_t dim2) {
  4036. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4037. auto maybe_layer = maybeCurrentDynamicLayer();
  4038. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4039. int64_t cur_level = maybe_layer->layerId();
  4040. if (!isBatchedAtLevel(grad_output, cur_level)) {
  4041. return at::_ops::diagonal_backward::call(grad_output, input_sizes, offset, dim1, dim2);
  4042. }
  4043. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  4044. auto results = batch_rule(grad_output_value, grad_output_bdim, input_sizes, offset, dim1, dim2);
  4045. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  4046. }
  4047. template <typename batch_rule_t, batch_rule_t batch_rule>
  4048. at::Tensor & fill_diagonal__generated_plumbing(at::Tensor & self, const at::Scalar & fill_value, bool wrap) {
  4049. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4050. auto maybe_layer = maybeCurrentDynamicLayer();
  4051. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  4052. int64_t cur_level = maybe_layer->layerId();
  4053. if (!isBatchedAtLevel(self, cur_level)) {
  4054. return at::_ops::fill_diagonal_::call(self, fill_value, wrap);
  4055. }
  4056. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4057. batch_rule(self_value, self_bdim, fill_value, wrap);
  4058. return self;
  4059. }
  4060. template <typename batch_rule_t, batch_rule_t batch_rule>
  4061. at::Tensor diff_generated_plumbing(const at::Tensor & self, int64_t n, int64_t dim, const ::std::optional<at::Tensor> & prepend, const ::std::optional<at::Tensor> & append) {
  4062. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4063. auto maybe_layer = maybeCurrentDynamicLayer();
  4064. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4065. int64_t cur_level = maybe_layer->layerId();
  4066. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(prepend, cur_level) && !isBatchedAtLevel(append, cur_level)) {
  4067. return at::_ops::diff::call(self, n, dim, prepend, append);
  4068. }
  4069. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4070. std::optional<Tensor> prepend_value;
  4071. std::optional<int64_t> prepend_bdim;
  4072. if (prepend) {
  4073. std::tie(prepend_value, prepend_bdim) = unwrapTensorAtLevel(prepend.value(), cur_level);
  4074. }
  4075. std::optional<Tensor> append_value;
  4076. std::optional<int64_t> append_bdim;
  4077. if (append) {
  4078. std::tie(append_value, append_bdim) = unwrapTensorAtLevel(append.value(), cur_level);
  4079. }
  4080. auto results = batch_rule(self_value, self_bdim, n, dim, prepend_value, prepend_bdim, append_value, append_bdim);
  4081. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  4082. }
  4083. template <typename batch_rule_t, batch_rule_t batch_rule>
  4084. ::std::vector<at::Tensor> gradient_scalarint_generated_plumbing(const at::Tensor & self, const ::std::optional<at::Scalar> & spacing, ::std::optional<int64_t> dim, int64_t edge_order) {
  4085. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4086. auto maybe_layer = maybeCurrentDynamicLayer();
  4087. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4088. int64_t cur_level = maybe_layer->layerId();
  4089. if (!isBatchedAtLevel(self, cur_level)) {
  4090. return at::_ops::gradient_scalarint::call(self, spacing, dim, edge_order);
  4091. }
  4092. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4093. auto results = batch_rule(self_value, self_bdim, spacing, dim, edge_order);
  4094. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  4095. }
  4096. template <typename batch_rule_t, batch_rule_t batch_rule>
  4097. ::std::vector<at::Tensor> gradient_scalararray_generated_plumbing(const at::Tensor & self, const at::Scalar & spacing, at::IntArrayRef dim, int64_t edge_order) {
  4098. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4099. auto maybe_layer = maybeCurrentDynamicLayer();
  4100. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4101. int64_t cur_level = maybe_layer->layerId();
  4102. if (!isBatchedAtLevel(self, cur_level)) {
  4103. return at::_ops::gradient_scalararray::call(self, spacing, dim, edge_order);
  4104. }
  4105. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4106. auto results = batch_rule(self_value, self_bdim, spacing, dim, edge_order);
  4107. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  4108. }
  4109. template <typename batch_rule_t, batch_rule_t batch_rule>
  4110. ::std::vector<at::Tensor> gradient_array_generated_plumbing(const at::Tensor & self, at::IntArrayRef dim, int64_t edge_order) {
  4111. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4112. auto maybe_layer = maybeCurrentDynamicLayer();
  4113. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4114. int64_t cur_level = maybe_layer->layerId();
  4115. if (!isBatchedAtLevel(self, cur_level)) {
  4116. return at::_ops::gradient_array::call(self, dim, edge_order);
  4117. }
  4118. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4119. auto results = batch_rule(self_value, self_bdim, dim, edge_order);
  4120. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  4121. }
  4122. template <typename batch_rule_t, batch_rule_t batch_rule>
  4123. ::std::vector<at::Tensor> gradient_scalarrayint_generated_plumbing(const at::Tensor & self, at::ArrayRef<at::Scalar> spacing, ::std::optional<int64_t> dim, int64_t edge_order) {
  4124. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4125. auto maybe_layer = maybeCurrentDynamicLayer();
  4126. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4127. int64_t cur_level = maybe_layer->layerId();
  4128. if (!isBatchedAtLevel(self, cur_level)) {
  4129. return at::_ops::gradient_scalarrayint::call(self, spacing, dim, edge_order);
  4130. }
  4131. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4132. auto results = batch_rule(self_value, self_bdim, spacing, dim, edge_order);
  4133. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  4134. }
  4135. template <typename batch_rule_t, batch_rule_t batch_rule>
  4136. ::std::vector<at::Tensor> gradient_scalarrayarray_generated_plumbing(const at::Tensor & self, at::ArrayRef<at::Scalar> spacing, at::IntArrayRef dim, int64_t edge_order) {
  4137. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4138. auto maybe_layer = maybeCurrentDynamicLayer();
  4139. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4140. int64_t cur_level = maybe_layer->layerId();
  4141. if (!isBatchedAtLevel(self, cur_level)) {
  4142. return at::_ops::gradient_scalarrayarray::call(self, spacing, dim, edge_order);
  4143. }
  4144. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4145. auto results = batch_rule(self_value, self_bdim, spacing, dim, edge_order);
  4146. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  4147. }
  4148. template <typename batch_rule_t, batch_rule_t batch_rule>
  4149. ::std::vector<at::Tensor> gradient_tensorarrayint_generated_plumbing(const at::Tensor & self, at::TensorList spacing, ::std::optional<int64_t> dim, int64_t edge_order) {
  4150. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4151. auto maybe_layer = maybeCurrentDynamicLayer();
  4152. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4153. int64_t cur_level = maybe_layer->layerId();
  4154. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(spacing, cur_level)) {
  4155. return at::_ops::gradient_tensorarrayint::call(self, spacing, dim, edge_order);
  4156. }
  4157. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4158. auto results = batch_rule(self_value, self_bdim, spacing, dim, edge_order);
  4159. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  4160. }
  4161. template <typename batch_rule_t, batch_rule_t batch_rule>
  4162. ::std::vector<at::Tensor> gradient_tensorarray_generated_plumbing(const at::Tensor & self, at::TensorList spacing, at::IntArrayRef dim, int64_t edge_order) {
  4163. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4164. auto maybe_layer = maybeCurrentDynamicLayer();
  4165. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4166. int64_t cur_level = maybe_layer->layerId();
  4167. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(spacing, cur_level)) {
  4168. return at::_ops::gradient_tensorarray::call(self, spacing, dim, edge_order);
  4169. }
  4170. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4171. auto results = batch_rule(self_value, self_bdim, spacing, dim, edge_order);
  4172. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  4173. }
  4174. template <typename batch_rule_t, batch_rule_t batch_rule>
  4175. at::Tensor div_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  4176. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4177. auto maybe_layer = maybeCurrentDynamicLayer();
  4178. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4179. int64_t cur_level = maybe_layer->layerId();
  4180. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  4181. return at::_ops::div_Tensor::call(self, other);
  4182. }
  4183. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4184. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  4185. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  4186. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  4187. }
  4188. template <typename batch_rule_t, batch_rule_t batch_rule>
  4189. at::Tensor & div__Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  4190. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4191. auto maybe_layer = maybeCurrentDynamicLayer();
  4192. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  4193. int64_t cur_level = maybe_layer->layerId();
  4194. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  4195. return at::_ops::div__Tensor::call(self, other);
  4196. }
  4197. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4198. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  4199. batch_rule(self_value, self_bdim, other_value, other_bdim);
  4200. return self;
  4201. }
  4202. template <typename batch_rule_t, batch_rule_t batch_rule>
  4203. at::Tensor div_Tensor_mode_generated_plumbing(const at::Tensor & self, const at::Tensor & other, ::std::optional<c10::string_view> rounding_mode) {
  4204. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4205. auto maybe_layer = maybeCurrentDynamicLayer();
  4206. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4207. int64_t cur_level = maybe_layer->layerId();
  4208. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  4209. return at::_ops::div_Tensor_mode::call(self, other, rounding_mode);
  4210. }
  4211. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4212. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  4213. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim, rounding_mode);
  4214. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  4215. }
  4216. template <typename batch_rule_t, batch_rule_t batch_rule>
  4217. at::Tensor & div__Tensor_mode_generated_plumbing(at::Tensor & self, const at::Tensor & other, ::std::optional<c10::string_view> rounding_mode) {
  4218. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4219. auto maybe_layer = maybeCurrentDynamicLayer();
  4220. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  4221. int64_t cur_level = maybe_layer->layerId();
  4222. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  4223. return at::_ops::div__Tensor_mode::call(self, other, rounding_mode);
  4224. }
  4225. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4226. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  4227. batch_rule(self_value, self_bdim, other_value, other_bdim, rounding_mode);
  4228. return self;
  4229. }
  4230. template <typename batch_rule_t, batch_rule_t batch_rule>
  4231. at::Tensor div_Scalar_generated_plumbing(const at::Tensor & self, const at::Scalar & other) {
  4232. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4233. auto maybe_layer = maybeCurrentDynamicLayer();
  4234. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4235. int64_t cur_level = maybe_layer->layerId();
  4236. if (!isBatchedAtLevel(self, cur_level)) {
  4237. return at::_ops::div_Scalar::call(self, other);
  4238. }
  4239. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4240. auto results = batch_rule(self_value, self_bdim, other);
  4241. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  4242. }
  4243. template <typename batch_rule_t, batch_rule_t batch_rule>
  4244. at::Tensor & div__Scalar_generated_plumbing(at::Tensor & self, const at::Scalar & other) {
  4245. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4246. auto maybe_layer = maybeCurrentDynamicLayer();
  4247. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  4248. int64_t cur_level = maybe_layer->layerId();
  4249. if (!isBatchedAtLevel(self, cur_level)) {
  4250. return at::_ops::div__Scalar::call(self, other);
  4251. }
  4252. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4253. batch_rule(self_value, self_bdim, other);
  4254. return self;
  4255. }
  4256. template <typename batch_rule_t, batch_rule_t batch_rule>
  4257. at::Tensor div_Scalar_mode_generated_plumbing(const at::Tensor & self, const at::Scalar & other, ::std::optional<c10::string_view> rounding_mode) {
  4258. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4259. auto maybe_layer = maybeCurrentDynamicLayer();
  4260. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4261. int64_t cur_level = maybe_layer->layerId();
  4262. if (!isBatchedAtLevel(self, cur_level)) {
  4263. return at::_ops::div_Scalar_mode::call(self, other, rounding_mode);
  4264. }
  4265. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4266. auto results = batch_rule(self_value, self_bdim, other, rounding_mode);
  4267. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  4268. }
  4269. template <typename batch_rule_t, batch_rule_t batch_rule>
  4270. at::Tensor & div__Scalar_mode_generated_plumbing(at::Tensor & self, const at::Scalar & other, ::std::optional<c10::string_view> rounding_mode) {
  4271. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4272. auto maybe_layer = maybeCurrentDynamicLayer();
  4273. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  4274. int64_t cur_level = maybe_layer->layerId();
  4275. if (!isBatchedAtLevel(self, cur_level)) {
  4276. return at::_ops::div__Scalar_mode::call(self, other, rounding_mode);
  4277. }
  4278. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4279. batch_rule(self_value, self_bdim, other, rounding_mode);
  4280. return self;
  4281. }
  4282. template <typename batch_rule_t, batch_rule_t batch_rule>
  4283. at::Tensor divide_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  4284. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4285. auto maybe_layer = maybeCurrentDynamicLayer();
  4286. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4287. int64_t cur_level = maybe_layer->layerId();
  4288. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  4289. return at::_ops::divide_Tensor::call(self, other);
  4290. }
  4291. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4292. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  4293. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  4294. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  4295. }
  4296. template <typename batch_rule_t, batch_rule_t batch_rule>
  4297. at::Tensor & divide__Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  4298. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4299. auto maybe_layer = maybeCurrentDynamicLayer();
  4300. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  4301. int64_t cur_level = maybe_layer->layerId();
  4302. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  4303. return at::_ops::divide__Tensor::call(self, other);
  4304. }
  4305. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4306. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  4307. batch_rule(self_value, self_bdim, other_value, other_bdim);
  4308. return self;
  4309. }
  4310. template <typename batch_rule_t, batch_rule_t batch_rule>
  4311. at::Tensor divide_Scalar_generated_plumbing(const at::Tensor & self, const at::Scalar & other) {
  4312. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4313. auto maybe_layer = maybeCurrentDynamicLayer();
  4314. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4315. int64_t cur_level = maybe_layer->layerId();
  4316. if (!isBatchedAtLevel(self, cur_level)) {
  4317. return at::_ops::divide_Scalar::call(self, other);
  4318. }
  4319. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4320. auto results = batch_rule(self_value, self_bdim, other);
  4321. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  4322. }
  4323. template <typename batch_rule_t, batch_rule_t batch_rule>
  4324. at::Tensor & divide__Scalar_generated_plumbing(at::Tensor & self, const at::Scalar & other) {
  4325. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4326. auto maybe_layer = maybeCurrentDynamicLayer();
  4327. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  4328. int64_t cur_level = maybe_layer->layerId();
  4329. if (!isBatchedAtLevel(self, cur_level)) {
  4330. return at::_ops::divide__Scalar::call(self, other);
  4331. }
  4332. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4333. batch_rule(self_value, self_bdim, other);
  4334. return self;
  4335. }
  4336. template <typename batch_rule_t, batch_rule_t batch_rule>
  4337. at::Tensor divide_Tensor_mode_generated_plumbing(const at::Tensor & self, const at::Tensor & other, ::std::optional<c10::string_view> rounding_mode) {
  4338. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4339. auto maybe_layer = maybeCurrentDynamicLayer();
  4340. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4341. int64_t cur_level = maybe_layer->layerId();
  4342. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  4343. return at::_ops::divide_Tensor_mode::call(self, other, rounding_mode);
  4344. }
  4345. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4346. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  4347. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim, rounding_mode);
  4348. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  4349. }
  4350. template <typename batch_rule_t, batch_rule_t batch_rule>
  4351. at::Tensor & divide__Tensor_mode_generated_plumbing(at::Tensor & self, const at::Tensor & other, ::std::optional<c10::string_view> rounding_mode) {
  4352. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4353. auto maybe_layer = maybeCurrentDynamicLayer();
  4354. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  4355. int64_t cur_level = maybe_layer->layerId();
  4356. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  4357. return at::_ops::divide__Tensor_mode::call(self, other, rounding_mode);
  4358. }
  4359. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4360. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  4361. batch_rule(self_value, self_bdim, other_value, other_bdim, rounding_mode);
  4362. return self;
  4363. }
  4364. template <typename batch_rule_t, batch_rule_t batch_rule>
  4365. at::Tensor divide_Scalar_mode_generated_plumbing(const at::Tensor & self, const at::Scalar & other, ::std::optional<c10::string_view> rounding_mode) {
  4366. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4367. auto maybe_layer = maybeCurrentDynamicLayer();
  4368. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4369. int64_t cur_level = maybe_layer->layerId();
  4370. if (!isBatchedAtLevel(self, cur_level)) {
  4371. return at::_ops::divide_Scalar_mode::call(self, other, rounding_mode);
  4372. }
  4373. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4374. auto results = batch_rule(self_value, self_bdim, other, rounding_mode);
  4375. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  4376. }
  4377. template <typename batch_rule_t, batch_rule_t batch_rule>
  4378. at::Tensor & divide__Scalar_mode_generated_plumbing(at::Tensor & self, const at::Scalar & other, ::std::optional<c10::string_view> rounding_mode) {
  4379. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4380. auto maybe_layer = maybeCurrentDynamicLayer();
  4381. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  4382. int64_t cur_level = maybe_layer->layerId();
  4383. if (!isBatchedAtLevel(self, cur_level)) {
  4384. return at::_ops::divide__Scalar_mode::call(self, other, rounding_mode);
  4385. }
  4386. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4387. batch_rule(self_value, self_bdim, other, rounding_mode);
  4388. return self;
  4389. }
  4390. template <typename batch_rule_t, batch_rule_t batch_rule>
  4391. at::Tensor true_divide_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  4392. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4393. auto maybe_layer = maybeCurrentDynamicLayer();
  4394. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4395. int64_t cur_level = maybe_layer->layerId();
  4396. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  4397. return at::_ops::true_divide_Tensor::call(self, other);
  4398. }
  4399. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4400. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  4401. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  4402. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  4403. }
  4404. template <typename batch_rule_t, batch_rule_t batch_rule>
  4405. at::Tensor & true_divide__Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  4406. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4407. auto maybe_layer = maybeCurrentDynamicLayer();
  4408. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  4409. int64_t cur_level = maybe_layer->layerId();
  4410. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  4411. return at::_ops::true_divide__Tensor::call(self, other);
  4412. }
  4413. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4414. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  4415. batch_rule(self_value, self_bdim, other_value, other_bdim);
  4416. return self;
  4417. }
  4418. template <typename batch_rule_t, batch_rule_t batch_rule>
  4419. at::Tensor true_divide_Scalar_generated_plumbing(const at::Tensor & self, const at::Scalar & other) {
  4420. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4421. auto maybe_layer = maybeCurrentDynamicLayer();
  4422. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4423. int64_t cur_level = maybe_layer->layerId();
  4424. if (!isBatchedAtLevel(self, cur_level)) {
  4425. return at::_ops::true_divide_Scalar::call(self, other);
  4426. }
  4427. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4428. auto results = batch_rule(self_value, self_bdim, other);
  4429. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  4430. }
  4431. template <typename batch_rule_t, batch_rule_t batch_rule>
  4432. at::Tensor & true_divide__Scalar_generated_plumbing(at::Tensor & self, const at::Scalar & other) {
  4433. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4434. auto maybe_layer = maybeCurrentDynamicLayer();
  4435. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  4436. int64_t cur_level = maybe_layer->layerId();
  4437. if (!isBatchedAtLevel(self, cur_level)) {
  4438. return at::_ops::true_divide__Scalar::call(self, other);
  4439. }
  4440. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4441. batch_rule(self_value, self_bdim, other);
  4442. return self;
  4443. }
  4444. template <typename batch_rule_t, batch_rule_t batch_rule>
  4445. at::Tensor dot_generated_plumbing(const at::Tensor & self, const at::Tensor & tensor) {
  4446. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4447. auto maybe_layer = maybeCurrentDynamicLayer();
  4448. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4449. int64_t cur_level = maybe_layer->layerId();
  4450. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(tensor, cur_level)) {
  4451. return at::_ops::dot::call(self, tensor);
  4452. }
  4453. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4454. auto [tensor_value, tensor_bdim] = unwrapTensorAtLevel(tensor, cur_level);
  4455. auto results = batch_rule(self_value, self_bdim, tensor_value, tensor_bdim);
  4456. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  4457. }
  4458. template <typename batch_rule_t, batch_rule_t batch_rule>
  4459. at::Tensor vdot_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  4460. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4461. auto maybe_layer = maybeCurrentDynamicLayer();
  4462. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4463. int64_t cur_level = maybe_layer->layerId();
  4464. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  4465. return at::_ops::vdot::call(self, other);
  4466. }
  4467. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4468. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  4469. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  4470. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  4471. }
  4472. template <typename batch_rule_t, batch_rule_t batch_rule>
  4473. at::Tensor einsum_generated_plumbing(c10::string_view equation, at::TensorList tensors, at::OptionalIntArrayRef path) {
  4474. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4475. auto maybe_layer = maybeCurrentDynamicLayer();
  4476. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4477. int64_t cur_level = maybe_layer->layerId();
  4478. if (!isBatchedAtLevel(tensors, cur_level)) {
  4479. return at::_ops::einsum::call(equation, tensors, path);
  4480. }
  4481. auto results = batch_rule(equation, tensors, path);
  4482. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  4483. }
  4484. template <typename batch_rule_t, batch_rule_t batch_rule>
  4485. at::Tensor embedding_generated_plumbing(const at::Tensor & weight, const at::Tensor & indices, c10::SymInt padding_idx, bool scale_grad_by_freq, bool sparse) {
  4486. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4487. auto maybe_layer = maybeCurrentDynamicLayer();
  4488. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4489. int64_t cur_level = maybe_layer->layerId();
  4490. if (!isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(indices, cur_level)) {
  4491. return at::_ops::embedding::call(weight, indices, padding_idx, scale_grad_by_freq, sparse);
  4492. }
  4493. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  4494. auto [indices_value, indices_bdim] = unwrapTensorAtLevel(indices, cur_level);
  4495. auto results = batch_rule(weight_value, weight_bdim, indices_value, indices_bdim, padding_idx, scale_grad_by_freq, sparse);
  4496. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  4497. }
  4498. template <typename batch_rule_t, batch_rule_t batch_rule>
  4499. at::Tensor embedding_backward_generated_plumbing(const at::Tensor & grad, const at::Tensor & indices, c10::SymInt num_weights, c10::SymInt padding_idx, bool scale_grad_by_freq, bool sparse) {
  4500. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4501. auto maybe_layer = maybeCurrentDynamicLayer();
  4502. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4503. int64_t cur_level = maybe_layer->layerId();
  4504. if (!isBatchedAtLevel(grad, cur_level) && !isBatchedAtLevel(indices, cur_level)) {
  4505. return at::_ops::embedding_backward::call(grad, indices, num_weights, padding_idx, scale_grad_by_freq, sparse);
  4506. }
  4507. auto [grad_value, grad_bdim] = unwrapTensorAtLevel(grad, cur_level);
  4508. auto [indices_value, indices_bdim] = unwrapTensorAtLevel(indices, cur_level);
  4509. auto results = batch_rule(grad_value, grad_bdim, indices_value, indices_bdim, num_weights, padding_idx, scale_grad_by_freq, sparse);
  4510. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  4511. }
  4512. template <typename batch_rule_t, batch_rule_t batch_rule>
  4513. at::Tensor embedding_dense_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & indices, c10::SymInt num_weights, c10::SymInt padding_idx, bool scale_grad_by_freq) {
  4514. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4515. auto maybe_layer = maybeCurrentDynamicLayer();
  4516. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4517. int64_t cur_level = maybe_layer->layerId();
  4518. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(indices, cur_level)) {
  4519. return at::_ops::embedding_dense_backward::call(grad_output, indices, num_weights, padding_idx, scale_grad_by_freq);
  4520. }
  4521. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  4522. auto [indices_value, indices_bdim] = unwrapTensorAtLevel(indices, cur_level);
  4523. auto results = batch_rule(grad_output_value, grad_output_bdim, indices_value, indices_bdim, num_weights, padding_idx, scale_grad_by_freq);
  4524. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  4525. }
  4526. template <typename batch_rule_t, batch_rule_t batch_rule>
  4527. at::Tensor & embedding_renorm__generated_plumbing(at::Tensor & self, const at::Tensor & indices, double max_norm, double norm_type) {
  4528. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4529. auto maybe_layer = maybeCurrentDynamicLayer();
  4530. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  4531. int64_t cur_level = maybe_layer->layerId();
  4532. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(indices, cur_level)) {
  4533. return at::_ops::embedding_renorm_::call(self, indices, max_norm, norm_type);
  4534. }
  4535. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4536. auto [indices_value, indices_bdim] = unwrapTensorAtLevel(indices, cur_level);
  4537. batch_rule(self_value, self_bdim, indices_value, indices_bdim, max_norm, norm_type);
  4538. return self;
  4539. }
  4540. template <typename batch_rule_t, batch_rule_t batch_rule>
  4541. at::Tensor embedding_sparse_backward_generated_plumbing(const at::Tensor & grad, const at::Tensor & indices, int64_t num_weights, int64_t padding_idx, bool scale_grad_by_freq) {
  4542. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4543. auto maybe_layer = maybeCurrentDynamicLayer();
  4544. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4545. int64_t cur_level = maybe_layer->layerId();
  4546. if (!isBatchedAtLevel(grad, cur_level) && !isBatchedAtLevel(indices, cur_level)) {
  4547. return at::_ops::embedding_sparse_backward::call(grad, indices, num_weights, padding_idx, scale_grad_by_freq);
  4548. }
  4549. auto [grad_value, grad_bdim] = unwrapTensorAtLevel(grad, cur_level);
  4550. auto [indices_value, indices_bdim] = unwrapTensorAtLevel(indices, cur_level);
  4551. auto results = batch_rule(grad_value, grad_bdim, indices_value, indices_bdim, num_weights, padding_idx, scale_grad_by_freq);
  4552. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  4553. }
  4554. template <typename batch_rule_t, batch_rule_t batch_rule>
  4555. ::std::tuple<at::Tensor,at::Tensor,at::Tensor,at::Tensor> _embedding_bag_forward_only_generated_plumbing(const at::Tensor & weight, const at::Tensor & indices, const at::Tensor & offsets, bool scale_grad_by_freq, int64_t mode, bool sparse, const ::std::optional<at::Tensor> & per_sample_weights, bool include_last_offset, int64_t padding_idx) {
  4556. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4557. auto maybe_layer = maybeCurrentDynamicLayer();
  4558. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4559. int64_t cur_level = maybe_layer->layerId();
  4560. if (!isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(indices, cur_level) && !isBatchedAtLevel(offsets, cur_level) && !isBatchedAtLevel(per_sample_weights, cur_level)) {
  4561. return at::_ops::_embedding_bag_forward_only::call(weight, indices, offsets, scale_grad_by_freq, mode, sparse, per_sample_weights, include_last_offset, padding_idx);
  4562. }
  4563. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  4564. auto [indices_value, indices_bdim] = unwrapTensorAtLevel(indices, cur_level);
  4565. auto [offsets_value, offsets_bdim] = unwrapTensorAtLevel(offsets, cur_level);
  4566. std::optional<Tensor> per_sample_weights_value;
  4567. std::optional<int64_t> per_sample_weights_bdim;
  4568. if (per_sample_weights) {
  4569. std::tie(per_sample_weights_value, per_sample_weights_bdim) = unwrapTensorAtLevel(per_sample_weights.value(), cur_level);
  4570. }
  4571. auto results = batch_rule(weight_value, weight_bdim, indices_value, indices_bdim, offsets_value, offsets_bdim, scale_grad_by_freq, mode, sparse, per_sample_weights_value, per_sample_weights_bdim, include_last_offset, padding_idx);
  4572. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level), makeBatched(std::get<6>(results), std::get<7>(results), cur_level));
  4573. }
  4574. template <typename batch_rule_t, batch_rule_t batch_rule>
  4575. ::std::tuple<at::Tensor,at::Tensor> _rowwise_prune_generated_plumbing(const at::Tensor & weight, const at::Tensor & mask, at::ScalarType compressed_indices_dtype) {
  4576. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4577. auto maybe_layer = maybeCurrentDynamicLayer();
  4578. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4579. int64_t cur_level = maybe_layer->layerId();
  4580. if (!isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(mask, cur_level)) {
  4581. return at::_ops::_rowwise_prune::call(weight, mask, compressed_indices_dtype);
  4582. }
  4583. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  4584. auto [mask_value, mask_bdim] = unwrapTensorAtLevel(mask, cur_level);
  4585. auto results = batch_rule(weight_value, weight_bdim, mask_value, mask_bdim, compressed_indices_dtype);
  4586. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  4587. }
  4588. template <typename batch_rule_t, batch_rule_t batch_rule>
  4589. at::Tensor row_stack_generated_plumbing(at::TensorList tensors) {
  4590. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4591. auto maybe_layer = maybeCurrentDynamicLayer();
  4592. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4593. int64_t cur_level = maybe_layer->layerId();
  4594. if (!isBatchedAtLevel(tensors, cur_level)) {
  4595. return at::_ops::row_stack::call(tensors);
  4596. }
  4597. auto results = batch_rule(tensors);
  4598. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  4599. }
  4600. template <typename batch_rule_t, batch_rule_t batch_rule>
  4601. ::std::tuple<at::Tensor,at::Tensor,at::Tensor,at::Tensor> embedding_bag_generated_plumbing(const at::Tensor & weight, const at::Tensor & indices, const at::Tensor & offsets, bool scale_grad_by_freq, int64_t mode, bool sparse, const ::std::optional<at::Tensor> & per_sample_weights, bool include_last_offset) {
  4602. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4603. auto maybe_layer = maybeCurrentDynamicLayer();
  4604. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4605. int64_t cur_level = maybe_layer->layerId();
  4606. if (!isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(indices, cur_level) && !isBatchedAtLevel(offsets, cur_level) && !isBatchedAtLevel(per_sample_weights, cur_level)) {
  4607. return at::_ops::embedding_bag::call(weight, indices, offsets, scale_grad_by_freq, mode, sparse, per_sample_weights, include_last_offset);
  4608. }
  4609. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  4610. auto [indices_value, indices_bdim] = unwrapTensorAtLevel(indices, cur_level);
  4611. auto [offsets_value, offsets_bdim] = unwrapTensorAtLevel(offsets, cur_level);
  4612. std::optional<Tensor> per_sample_weights_value;
  4613. std::optional<int64_t> per_sample_weights_bdim;
  4614. if (per_sample_weights) {
  4615. std::tie(per_sample_weights_value, per_sample_weights_bdim) = unwrapTensorAtLevel(per_sample_weights.value(), cur_level);
  4616. }
  4617. auto results = batch_rule(weight_value, weight_bdim, indices_value, indices_bdim, offsets_value, offsets_bdim, scale_grad_by_freq, mode, sparse, per_sample_weights_value, per_sample_weights_bdim, include_last_offset);
  4618. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level), makeBatched(std::get<6>(results), std::get<7>(results), cur_level));
  4619. }
  4620. template <typename batch_rule_t, batch_rule_t batch_rule>
  4621. ::std::tuple<at::Tensor,at::Tensor,at::Tensor,at::Tensor> embedding_bag_padding_idx_generated_plumbing(const at::Tensor & weight, const at::Tensor & indices, const at::Tensor & offsets, bool scale_grad_by_freq, int64_t mode, bool sparse, const ::std::optional<at::Tensor> & per_sample_weights, bool include_last_offset, ::std::optional<int64_t> padding_idx) {
  4622. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4623. auto maybe_layer = maybeCurrentDynamicLayer();
  4624. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4625. int64_t cur_level = maybe_layer->layerId();
  4626. if (!isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(indices, cur_level) && !isBatchedAtLevel(offsets, cur_level) && !isBatchedAtLevel(per_sample_weights, cur_level)) {
  4627. return at::_ops::embedding_bag_padding_idx::call(weight, indices, offsets, scale_grad_by_freq, mode, sparse, per_sample_weights, include_last_offset, padding_idx);
  4628. }
  4629. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  4630. auto [indices_value, indices_bdim] = unwrapTensorAtLevel(indices, cur_level);
  4631. auto [offsets_value, offsets_bdim] = unwrapTensorAtLevel(offsets, cur_level);
  4632. std::optional<Tensor> per_sample_weights_value;
  4633. std::optional<int64_t> per_sample_weights_bdim;
  4634. if (per_sample_weights) {
  4635. std::tie(per_sample_weights_value, per_sample_weights_bdim) = unwrapTensorAtLevel(per_sample_weights.value(), cur_level);
  4636. }
  4637. auto results = batch_rule(weight_value, weight_bdim, indices_value, indices_bdim, offsets_value, offsets_bdim, scale_grad_by_freq, mode, sparse, per_sample_weights_value, per_sample_weights_bdim, include_last_offset, padding_idx);
  4638. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level), makeBatched(std::get<6>(results), std::get<7>(results), cur_level));
  4639. }
  4640. template <typename batch_rule_t, batch_rule_t batch_rule>
  4641. ::std::tuple<at::Tensor,at::Tensor,at::Tensor,at::Tensor> _embedding_bag_generated_plumbing(const at::Tensor & weight, const at::Tensor & indices, const at::Tensor & offsets, bool scale_grad_by_freq, int64_t mode, bool sparse, const ::std::optional<at::Tensor> & per_sample_weights, bool include_last_offset, int64_t padding_idx) {
  4642. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4643. auto maybe_layer = maybeCurrentDynamicLayer();
  4644. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4645. int64_t cur_level = maybe_layer->layerId();
  4646. if (!isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(indices, cur_level) && !isBatchedAtLevel(offsets, cur_level) && !isBatchedAtLevel(per_sample_weights, cur_level)) {
  4647. return at::_ops::_embedding_bag::call(weight, indices, offsets, scale_grad_by_freq, mode, sparse, per_sample_weights, include_last_offset, padding_idx);
  4648. }
  4649. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  4650. auto [indices_value, indices_bdim] = unwrapTensorAtLevel(indices, cur_level);
  4651. auto [offsets_value, offsets_bdim] = unwrapTensorAtLevel(offsets, cur_level);
  4652. std::optional<Tensor> per_sample_weights_value;
  4653. std::optional<int64_t> per_sample_weights_bdim;
  4654. if (per_sample_weights) {
  4655. std::tie(per_sample_weights_value, per_sample_weights_bdim) = unwrapTensorAtLevel(per_sample_weights.value(), cur_level);
  4656. }
  4657. auto results = batch_rule(weight_value, weight_bdim, indices_value, indices_bdim, offsets_value, offsets_bdim, scale_grad_by_freq, mode, sparse, per_sample_weights_value, per_sample_weights_bdim, include_last_offset, padding_idx);
  4658. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level), makeBatched(std::get<6>(results), std::get<7>(results), cur_level));
  4659. }
  4660. template <typename batch_rule_t, batch_rule_t batch_rule>
  4661. at::Tensor _embedding_bag_backward_generated_plumbing(const at::Tensor & grad, const at::Tensor & indices, const at::Tensor & offsets, const at::Tensor & offset2bag, const at::Tensor & bag_size, const at::Tensor & maximum_indices, c10::SymInt num_weights, bool scale_grad_by_freq, int64_t mode, bool sparse, const ::std::optional<at::Tensor> & per_sample_weights, int64_t padding_idx) {
  4662. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4663. auto maybe_layer = maybeCurrentDynamicLayer();
  4664. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4665. int64_t cur_level = maybe_layer->layerId();
  4666. if (!isBatchedAtLevel(grad, cur_level) && !isBatchedAtLevel(indices, cur_level) && !isBatchedAtLevel(offsets, cur_level) && !isBatchedAtLevel(offset2bag, cur_level) && !isBatchedAtLevel(bag_size, cur_level) && !isBatchedAtLevel(maximum_indices, cur_level) && !isBatchedAtLevel(per_sample_weights, cur_level)) {
  4667. return at::_ops::_embedding_bag_backward::call(grad, indices, offsets, offset2bag, bag_size, maximum_indices, num_weights, scale_grad_by_freq, mode, sparse, per_sample_weights, padding_idx);
  4668. }
  4669. auto [grad_value, grad_bdim] = unwrapTensorAtLevel(grad, cur_level);
  4670. auto [indices_value, indices_bdim] = unwrapTensorAtLevel(indices, cur_level);
  4671. auto [offsets_value, offsets_bdim] = unwrapTensorAtLevel(offsets, cur_level);
  4672. auto [offset2bag_value, offset2bag_bdim] = unwrapTensorAtLevel(offset2bag, cur_level);
  4673. auto [bag_size_value, bag_size_bdim] = unwrapTensorAtLevel(bag_size, cur_level);
  4674. auto [maximum_indices_value, maximum_indices_bdim] = unwrapTensorAtLevel(maximum_indices, cur_level);
  4675. std::optional<Tensor> per_sample_weights_value;
  4676. std::optional<int64_t> per_sample_weights_bdim;
  4677. if (per_sample_weights) {
  4678. std::tie(per_sample_weights_value, per_sample_weights_bdim) = unwrapTensorAtLevel(per_sample_weights.value(), cur_level);
  4679. }
  4680. auto results = batch_rule(grad_value, grad_bdim, indices_value, indices_bdim, offsets_value, offsets_bdim, offset2bag_value, offset2bag_bdim, bag_size_value, bag_size_bdim, maximum_indices_value, maximum_indices_bdim, num_weights, scale_grad_by_freq, mode, sparse, per_sample_weights_value, per_sample_weights_bdim, padding_idx);
  4681. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  4682. }
  4683. template <typename batch_rule_t, batch_rule_t batch_rule>
  4684. at::Tensor _embedding_bag_sparse_backward_generated_plumbing(const at::Tensor & grad, const at::Tensor & indices, const at::Tensor & offsets, const at::Tensor & offset2bag, const at::Tensor & bag_size, c10::SymInt num_weights, bool scale_grad_by_freq, int64_t mode, const ::std::optional<at::Tensor> & per_sample_weights, int64_t padding_idx) {
  4685. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4686. auto maybe_layer = maybeCurrentDynamicLayer();
  4687. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4688. int64_t cur_level = maybe_layer->layerId();
  4689. if (!isBatchedAtLevel(grad, cur_level) && !isBatchedAtLevel(indices, cur_level) && !isBatchedAtLevel(offsets, cur_level) && !isBatchedAtLevel(offset2bag, cur_level) && !isBatchedAtLevel(bag_size, cur_level) && !isBatchedAtLevel(per_sample_weights, cur_level)) {
  4690. return at::_ops::_embedding_bag_sparse_backward::call(grad, indices, offsets, offset2bag, bag_size, num_weights, scale_grad_by_freq, mode, per_sample_weights, padding_idx);
  4691. }
  4692. auto [grad_value, grad_bdim] = unwrapTensorAtLevel(grad, cur_level);
  4693. auto [indices_value, indices_bdim] = unwrapTensorAtLevel(indices, cur_level);
  4694. auto [offsets_value, offsets_bdim] = unwrapTensorAtLevel(offsets, cur_level);
  4695. auto [offset2bag_value, offset2bag_bdim] = unwrapTensorAtLevel(offset2bag, cur_level);
  4696. auto [bag_size_value, bag_size_bdim] = unwrapTensorAtLevel(bag_size, cur_level);
  4697. std::optional<Tensor> per_sample_weights_value;
  4698. std::optional<int64_t> per_sample_weights_bdim;
  4699. if (per_sample_weights) {
  4700. std::tie(per_sample_weights_value, per_sample_weights_bdim) = unwrapTensorAtLevel(per_sample_weights.value(), cur_level);
  4701. }
  4702. auto results = batch_rule(grad_value, grad_bdim, indices_value, indices_bdim, offsets_value, offsets_bdim, offset2bag_value, offset2bag_bdim, bag_size_value, bag_size_bdim, num_weights, scale_grad_by_freq, mode, per_sample_weights_value, per_sample_weights_bdim, padding_idx);
  4703. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  4704. }
  4705. template <typename batch_rule_t, batch_rule_t batch_rule>
  4706. at::Tensor _embedding_bag_dense_backward_generated_plumbing(const at::Tensor & grad, const at::Tensor & indices, const at::Tensor & offset2bag, const at::Tensor & bag_size, const at::Tensor & maximum_indices, c10::SymInt num_weights, bool scale_grad_by_freq, int64_t mode, const ::std::optional<at::Tensor> & per_sample_weights, int64_t padding_idx) {
  4707. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4708. auto maybe_layer = maybeCurrentDynamicLayer();
  4709. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4710. int64_t cur_level = maybe_layer->layerId();
  4711. if (!isBatchedAtLevel(grad, cur_level) && !isBatchedAtLevel(indices, cur_level) && !isBatchedAtLevel(offset2bag, cur_level) && !isBatchedAtLevel(bag_size, cur_level) && !isBatchedAtLevel(maximum_indices, cur_level) && !isBatchedAtLevel(per_sample_weights, cur_level)) {
  4712. return at::_ops::_embedding_bag_dense_backward::call(grad, indices, offset2bag, bag_size, maximum_indices, num_weights, scale_grad_by_freq, mode, per_sample_weights, padding_idx);
  4713. }
  4714. auto [grad_value, grad_bdim] = unwrapTensorAtLevel(grad, cur_level);
  4715. auto [indices_value, indices_bdim] = unwrapTensorAtLevel(indices, cur_level);
  4716. auto [offset2bag_value, offset2bag_bdim] = unwrapTensorAtLevel(offset2bag, cur_level);
  4717. auto [bag_size_value, bag_size_bdim] = unwrapTensorAtLevel(bag_size, cur_level);
  4718. auto [maximum_indices_value, maximum_indices_bdim] = unwrapTensorAtLevel(maximum_indices, cur_level);
  4719. std::optional<Tensor> per_sample_weights_value;
  4720. std::optional<int64_t> per_sample_weights_bdim;
  4721. if (per_sample_weights) {
  4722. std::tie(per_sample_weights_value, per_sample_weights_bdim) = unwrapTensorAtLevel(per_sample_weights.value(), cur_level);
  4723. }
  4724. auto results = batch_rule(grad_value, grad_bdim, indices_value, indices_bdim, offset2bag_value, offset2bag_bdim, bag_size_value, bag_size_bdim, maximum_indices_value, maximum_indices_bdim, num_weights, scale_grad_by_freq, mode, per_sample_weights_value, per_sample_weights_bdim, padding_idx);
  4725. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  4726. }
  4727. template <typename batch_rule_t, batch_rule_t batch_rule>
  4728. at::Tensor _embedding_bag_per_sample_weights_backward_generated_plumbing(const at::Tensor & grad, const at::Tensor & weight, const at::Tensor & indices, const at::Tensor & offsets, const at::Tensor & offset2bag, int64_t mode, int64_t padding_idx) {
  4729. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4730. auto maybe_layer = maybeCurrentDynamicLayer();
  4731. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4732. int64_t cur_level = maybe_layer->layerId();
  4733. if (!isBatchedAtLevel(grad, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(indices, cur_level) && !isBatchedAtLevel(offsets, cur_level) && !isBatchedAtLevel(offset2bag, cur_level)) {
  4734. return at::_ops::_embedding_bag_per_sample_weights_backward::call(grad, weight, indices, offsets, offset2bag, mode, padding_idx);
  4735. }
  4736. auto [grad_value, grad_bdim] = unwrapTensorAtLevel(grad, cur_level);
  4737. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  4738. auto [indices_value, indices_bdim] = unwrapTensorAtLevel(indices, cur_level);
  4739. auto [offsets_value, offsets_bdim] = unwrapTensorAtLevel(offsets, cur_level);
  4740. auto [offset2bag_value, offset2bag_bdim] = unwrapTensorAtLevel(offset2bag, cur_level);
  4741. auto results = batch_rule(grad_value, grad_bdim, weight_value, weight_bdim, indices_value, indices_bdim, offsets_value, offsets_bdim, offset2bag_value, offset2bag_bdim, mode, padding_idx);
  4742. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  4743. }
  4744. template <typename batch_rule_t, batch_rule_t batch_rule>
  4745. at::Tensor new_empty_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef size, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory) {
  4746. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4747. auto maybe_layer = maybeCurrentDynamicLayer();
  4748. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4749. int64_t cur_level = maybe_layer->layerId();
  4750. if (!isBatchedAtLevel(self, cur_level)) {
  4751. return at::_ops::new_empty::call(self, size, dtype, layout, device, pin_memory);
  4752. }
  4753. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4754. auto results = batch_rule(self_value, self_bdim, size, dtype, layout, device, pin_memory);
  4755. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  4756. }
  4757. template <typename batch_rule_t, batch_rule_t batch_rule>
  4758. at::Tensor new_empty_strided_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef size, c10::SymIntArrayRef stride, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory) {
  4759. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4760. auto maybe_layer = maybeCurrentDynamicLayer();
  4761. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4762. int64_t cur_level = maybe_layer->layerId();
  4763. if (!isBatchedAtLevel(self, cur_level)) {
  4764. return at::_ops::new_empty_strided::call(self, size, stride, dtype, layout, device, pin_memory);
  4765. }
  4766. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4767. auto results = batch_rule(self_value, self_bdim, size, stride, dtype, layout, device, pin_memory);
  4768. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  4769. }
  4770. template <typename batch_rule_t, batch_rule_t batch_rule>
  4771. at::Tensor new_full_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef size, const at::Scalar & fill_value, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory) {
  4772. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4773. auto maybe_layer = maybeCurrentDynamicLayer();
  4774. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4775. int64_t cur_level = maybe_layer->layerId();
  4776. if (!isBatchedAtLevel(self, cur_level)) {
  4777. return at::_ops::new_full::call(self, size, fill_value, dtype, layout, device, pin_memory);
  4778. }
  4779. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4780. auto results = batch_rule(self_value, self_bdim, size, fill_value, dtype, layout, device, pin_memory);
  4781. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  4782. }
  4783. template <typename batch_rule_t, batch_rule_t batch_rule>
  4784. at::Tensor new_zeros_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef size, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory) {
  4785. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4786. auto maybe_layer = maybeCurrentDynamicLayer();
  4787. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4788. int64_t cur_level = maybe_layer->layerId();
  4789. if (!isBatchedAtLevel(self, cur_level)) {
  4790. return at::_ops::new_zeros::call(self, size, dtype, layout, device, pin_memory);
  4791. }
  4792. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4793. auto results = batch_rule(self_value, self_bdim, size, dtype, layout, device, pin_memory);
  4794. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  4795. }
  4796. template <typename batch_rule_t, batch_rule_t batch_rule>
  4797. at::Tensor new_ones_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef size, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory) {
  4798. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4799. auto maybe_layer = maybeCurrentDynamicLayer();
  4800. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4801. int64_t cur_level = maybe_layer->layerId();
  4802. if (!isBatchedAtLevel(self, cur_level)) {
  4803. return at::_ops::new_ones::call(self, size, dtype, layout, device, pin_memory);
  4804. }
  4805. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4806. auto results = batch_rule(self_value, self_bdim, size, dtype, layout, device, pin_memory);
  4807. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  4808. }
  4809. template <typename batch_rule_t, batch_rule_t batch_rule>
  4810. at::Tensor _empty_per_channel_affine_quantized_generated_plumbing(c10::SymIntArrayRef size, const at::Tensor & scales, const at::Tensor & zero_points, int64_t axis, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory, ::std::optional<at::MemoryFormat> memory_format) {
  4811. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4812. auto maybe_layer = maybeCurrentDynamicLayer();
  4813. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4814. int64_t cur_level = maybe_layer->layerId();
  4815. if (!isBatchedAtLevel(scales, cur_level) && !isBatchedAtLevel(zero_points, cur_level)) {
  4816. return at::_ops::_empty_per_channel_affine_quantized::call(size, scales, zero_points, axis, dtype, layout, device, pin_memory, memory_format);
  4817. }
  4818. auto [scales_value, scales_bdim] = unwrapTensorAtLevel(scales, cur_level);
  4819. auto [zero_points_value, zero_points_bdim] = unwrapTensorAtLevel(zero_points, cur_level);
  4820. auto results = batch_rule(size, scales_value, scales_bdim, zero_points_value, zero_points_bdim, axis, dtype, layout, device, pin_memory, memory_format);
  4821. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  4822. }
  4823. template <typename batch_rule_t, batch_rule_t batch_rule>
  4824. const at::Tensor & _resize_output__generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef size, at::Device device) {
  4825. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4826. auto maybe_layer = maybeCurrentDynamicLayer();
  4827. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  4828. int64_t cur_level = maybe_layer->layerId();
  4829. if (!isBatchedAtLevel(self, cur_level)) {
  4830. return at::_ops::_resize_output_::call(self, size, device);
  4831. }
  4832. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4833. batch_rule(self_value, self_bdim, size, device);
  4834. return self;
  4835. }
  4836. template <typename batch_rule_t, batch_rule_t batch_rule>
  4837. at::Tensor empty_quantized_generated_plumbing(at::IntArrayRef size, const at::Tensor & qtensor, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory, ::std::optional<at::MemoryFormat> memory_format) {
  4838. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4839. auto maybe_layer = maybeCurrentDynamicLayer();
  4840. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4841. int64_t cur_level = maybe_layer->layerId();
  4842. if (!isBatchedAtLevel(qtensor, cur_level)) {
  4843. return at::_ops::empty_quantized::call(size, qtensor, dtype, layout, device, pin_memory, memory_format);
  4844. }
  4845. auto [qtensor_value, qtensor_bdim] = unwrapTensorAtLevel(qtensor, cur_level);
  4846. auto results = batch_rule(size, qtensor_value, qtensor_bdim, dtype, layout, device, pin_memory, memory_format);
  4847. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  4848. }
  4849. template <typename batch_rule_t, batch_rule_t batch_rule>
  4850. at::Tensor empty_like_generated_plumbing(const at::Tensor & self, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory, ::std::optional<at::MemoryFormat> memory_format) {
  4851. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4852. auto maybe_layer = maybeCurrentDynamicLayer();
  4853. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4854. int64_t cur_level = maybe_layer->layerId();
  4855. if (!isBatchedAtLevel(self, cur_level)) {
  4856. return at::_ops::empty_like::call(self, dtype, layout, device, pin_memory, memory_format);
  4857. }
  4858. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4859. auto results = batch_rule(self_value, self_bdim, dtype, layout, device, pin_memory, memory_format);
  4860. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  4861. }
  4862. template <typename batch_rule_t, batch_rule_t batch_rule>
  4863. at::Tensor erf_generated_plumbing(const at::Tensor & self) {
  4864. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4865. auto maybe_layer = maybeCurrentDynamicLayer();
  4866. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4867. int64_t cur_level = maybe_layer->layerId();
  4868. if (!isBatchedAtLevel(self, cur_level)) {
  4869. return at::_ops::erf::call(self);
  4870. }
  4871. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4872. auto results = batch_rule(self_value, self_bdim);
  4873. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  4874. }
  4875. template <typename batch_rule_t, batch_rule_t batch_rule>
  4876. at::Tensor & erf__generated_plumbing(at::Tensor & self) {
  4877. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4878. auto maybe_layer = maybeCurrentDynamicLayer();
  4879. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  4880. int64_t cur_level = maybe_layer->layerId();
  4881. if (!isBatchedAtLevel(self, cur_level)) {
  4882. return at::_ops::erf_::call(self);
  4883. }
  4884. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4885. batch_rule(self_value, self_bdim);
  4886. return self;
  4887. }
  4888. template <typename batch_rule_t, batch_rule_t batch_rule>
  4889. at::Tensor erfc_generated_plumbing(const at::Tensor & self) {
  4890. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4891. auto maybe_layer = maybeCurrentDynamicLayer();
  4892. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4893. int64_t cur_level = maybe_layer->layerId();
  4894. if (!isBatchedAtLevel(self, cur_level)) {
  4895. return at::_ops::erfc::call(self);
  4896. }
  4897. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4898. auto results = batch_rule(self_value, self_bdim);
  4899. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  4900. }
  4901. template <typename batch_rule_t, batch_rule_t batch_rule>
  4902. at::Tensor & erfc__generated_plumbing(at::Tensor & self) {
  4903. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4904. auto maybe_layer = maybeCurrentDynamicLayer();
  4905. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  4906. int64_t cur_level = maybe_layer->layerId();
  4907. if (!isBatchedAtLevel(self, cur_level)) {
  4908. return at::_ops::erfc_::call(self);
  4909. }
  4910. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4911. batch_rule(self_value, self_bdim);
  4912. return self;
  4913. }
  4914. template <typename batch_rule_t, batch_rule_t batch_rule>
  4915. at::Tensor exp_generated_plumbing(const at::Tensor & self) {
  4916. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4917. auto maybe_layer = maybeCurrentDynamicLayer();
  4918. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4919. int64_t cur_level = maybe_layer->layerId();
  4920. if (!isBatchedAtLevel(self, cur_level)) {
  4921. return at::_ops::exp::call(self);
  4922. }
  4923. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4924. auto results = batch_rule(self_value, self_bdim);
  4925. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  4926. }
  4927. template <typename batch_rule_t, batch_rule_t batch_rule>
  4928. at::Tensor & exp__generated_plumbing(at::Tensor & self) {
  4929. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4930. auto maybe_layer = maybeCurrentDynamicLayer();
  4931. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  4932. int64_t cur_level = maybe_layer->layerId();
  4933. if (!isBatchedAtLevel(self, cur_level)) {
  4934. return at::_ops::exp_::call(self);
  4935. }
  4936. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4937. batch_rule(self_value, self_bdim);
  4938. return self;
  4939. }
  4940. template <typename batch_rule_t, batch_rule_t batch_rule>
  4941. at::Tensor exp2_generated_plumbing(const at::Tensor & self) {
  4942. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4943. auto maybe_layer = maybeCurrentDynamicLayer();
  4944. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4945. int64_t cur_level = maybe_layer->layerId();
  4946. if (!isBatchedAtLevel(self, cur_level)) {
  4947. return at::_ops::exp2::call(self);
  4948. }
  4949. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4950. auto results = batch_rule(self_value, self_bdim);
  4951. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  4952. }
  4953. template <typename batch_rule_t, batch_rule_t batch_rule>
  4954. at::Tensor & exp2__generated_plumbing(at::Tensor & self) {
  4955. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4956. auto maybe_layer = maybeCurrentDynamicLayer();
  4957. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  4958. int64_t cur_level = maybe_layer->layerId();
  4959. if (!isBatchedAtLevel(self, cur_level)) {
  4960. return at::_ops::exp2_::call(self);
  4961. }
  4962. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4963. batch_rule(self_value, self_bdim);
  4964. return self;
  4965. }
  4966. template <typename batch_rule_t, batch_rule_t batch_rule>
  4967. at::Tensor expm1_generated_plumbing(const at::Tensor & self) {
  4968. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4969. auto maybe_layer = maybeCurrentDynamicLayer();
  4970. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4971. int64_t cur_level = maybe_layer->layerId();
  4972. if (!isBatchedAtLevel(self, cur_level)) {
  4973. return at::_ops::expm1::call(self);
  4974. }
  4975. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4976. auto results = batch_rule(self_value, self_bdim);
  4977. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  4978. }
  4979. template <typename batch_rule_t, batch_rule_t batch_rule>
  4980. at::Tensor & expm1__generated_plumbing(at::Tensor & self) {
  4981. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4982. auto maybe_layer = maybeCurrentDynamicLayer();
  4983. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  4984. int64_t cur_level = maybe_layer->layerId();
  4985. if (!isBatchedAtLevel(self, cur_level)) {
  4986. return at::_ops::expm1_::call(self);
  4987. }
  4988. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  4989. batch_rule(self_value, self_bdim);
  4990. return self;
  4991. }
  4992. template <typename batch_rule_t, batch_rule_t batch_rule>
  4993. at::Tensor expand_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef size, bool implicit) {
  4994. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  4995. auto maybe_layer = maybeCurrentDynamicLayer();
  4996. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  4997. int64_t cur_level = maybe_layer->layerId();
  4998. if (!isBatchedAtLevel(self, cur_level)) {
  4999. return at::_ops::expand::call(self, size, implicit);
  5000. }
  5001. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5002. auto results = batch_rule(self_value, self_bdim, size, implicit);
  5003. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5004. }
  5005. template <typename batch_rule_t, batch_rule_t batch_rule>
  5006. at::Tensor expand_as_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  5007. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5008. auto maybe_layer = maybeCurrentDynamicLayer();
  5009. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5010. int64_t cur_level = maybe_layer->layerId();
  5011. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  5012. return at::_ops::expand_as::call(self, other);
  5013. }
  5014. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5015. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  5016. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  5017. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5018. }
  5019. template <typename batch_rule_t, batch_rule_t batch_rule>
  5020. at::Tensor flatten_using_ints_generated_plumbing(const at::Tensor & self, int64_t start_dim, int64_t end_dim) {
  5021. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5022. auto maybe_layer = maybeCurrentDynamicLayer();
  5023. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5024. int64_t cur_level = maybe_layer->layerId();
  5025. if (!isBatchedAtLevel(self, cur_level)) {
  5026. return at::_ops::flatten_using_ints::call(self, start_dim, end_dim);
  5027. }
  5028. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5029. auto results = batch_rule(self_value, self_bdim, start_dim, end_dim);
  5030. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5031. }
  5032. template <typename batch_rule_t, batch_rule_t batch_rule>
  5033. at::Tensor flatten_named_out_dim_generated_plumbing(const at::Tensor & self, int64_t start_dim, int64_t end_dim, at::Dimname out_dim) {
  5034. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5035. auto maybe_layer = maybeCurrentDynamicLayer();
  5036. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5037. int64_t cur_level = maybe_layer->layerId();
  5038. if (!isBatchedAtLevel(self, cur_level)) {
  5039. return at::_ops::flatten_named_out_dim::call(self, start_dim, end_dim, out_dim);
  5040. }
  5041. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5042. auto results = batch_rule(self_value, self_bdim, start_dim, end_dim, out_dim);
  5043. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5044. }
  5045. template <typename batch_rule_t, batch_rule_t batch_rule>
  5046. at::Tensor flatten_using_names_generated_plumbing(const at::Tensor & self, at::Dimname start_dim, at::Dimname end_dim, at::Dimname out_dim) {
  5047. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5048. auto maybe_layer = maybeCurrentDynamicLayer();
  5049. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5050. int64_t cur_level = maybe_layer->layerId();
  5051. if (!isBatchedAtLevel(self, cur_level)) {
  5052. return at::_ops::flatten_using_names::call(self, start_dim, end_dim, out_dim);
  5053. }
  5054. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5055. auto results = batch_rule(self_value, self_bdim, start_dim, end_dim, out_dim);
  5056. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5057. }
  5058. template <typename batch_rule_t, batch_rule_t batch_rule>
  5059. at::Tensor flatten_DimnameList_generated_plumbing(const at::Tensor & self, at::DimnameList dims, at::Dimname out_dim) {
  5060. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5061. auto maybe_layer = maybeCurrentDynamicLayer();
  5062. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5063. int64_t cur_level = maybe_layer->layerId();
  5064. if (!isBatchedAtLevel(self, cur_level)) {
  5065. return at::_ops::flatten_DimnameList::call(self, dims, out_dim);
  5066. }
  5067. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5068. auto results = batch_rule(self_value, self_bdim, dims, out_dim);
  5069. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5070. }
  5071. template <typename batch_rule_t, batch_rule_t batch_rule>
  5072. at::Tensor unflatten_int_generated_plumbing(const at::Tensor & self, int64_t dim, c10::SymIntArrayRef sizes) {
  5073. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5074. auto maybe_layer = maybeCurrentDynamicLayer();
  5075. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5076. int64_t cur_level = maybe_layer->layerId();
  5077. if (!isBatchedAtLevel(self, cur_level)) {
  5078. return at::_ops::unflatten_int::call(self, dim, sizes);
  5079. }
  5080. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5081. auto results = batch_rule(self_value, self_bdim, dim, sizes);
  5082. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5083. }
  5084. template <typename batch_rule_t, batch_rule_t batch_rule>
  5085. at::Tensor unflatten_Dimname_generated_plumbing(const at::Tensor & self, at::Dimname dim, c10::SymIntArrayRef sizes, at::DimnameList names) {
  5086. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5087. auto maybe_layer = maybeCurrentDynamicLayer();
  5088. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5089. int64_t cur_level = maybe_layer->layerId();
  5090. if (!isBatchedAtLevel(self, cur_level)) {
  5091. return at::_ops::unflatten_Dimname::call(self, dim, sizes, names);
  5092. }
  5093. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5094. auto results = batch_rule(self_value, self_bdim, dim, sizes, names);
  5095. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5096. }
  5097. template <typename batch_rule_t, batch_rule_t batch_rule>
  5098. at::Tensor fill_Scalar_generated_plumbing(const at::Tensor & self, const at::Scalar & value) {
  5099. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5100. auto maybe_layer = maybeCurrentDynamicLayer();
  5101. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5102. int64_t cur_level = maybe_layer->layerId();
  5103. if (!isBatchedAtLevel(self, cur_level)) {
  5104. return at::_ops::fill_Scalar::call(self, value);
  5105. }
  5106. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5107. auto results = batch_rule(self_value, self_bdim, value);
  5108. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5109. }
  5110. template <typename batch_rule_t, batch_rule_t batch_rule>
  5111. at::Tensor fill_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & value) {
  5112. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5113. auto maybe_layer = maybeCurrentDynamicLayer();
  5114. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5115. int64_t cur_level = maybe_layer->layerId();
  5116. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(value, cur_level)) {
  5117. return at::_ops::fill_Tensor::call(self, value);
  5118. }
  5119. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5120. auto [value_value, value_bdim] = unwrapTensorAtLevel(value, cur_level);
  5121. auto results = batch_rule(self_value, self_bdim, value_value, value_bdim);
  5122. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5123. }
  5124. template <typename batch_rule_t, batch_rule_t batch_rule>
  5125. at::Tensor & fill__Scalar_generated_plumbing(at::Tensor & self, const at::Scalar & value) {
  5126. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5127. auto maybe_layer = maybeCurrentDynamicLayer();
  5128. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  5129. int64_t cur_level = maybe_layer->layerId();
  5130. if (!isBatchedAtLevel(self, cur_level)) {
  5131. return at::_ops::fill__Scalar::call(self, value);
  5132. }
  5133. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5134. batch_rule(self_value, self_bdim, value);
  5135. return self;
  5136. }
  5137. template <typename batch_rule_t, batch_rule_t batch_rule>
  5138. at::Tensor & fill__Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & value) {
  5139. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5140. auto maybe_layer = maybeCurrentDynamicLayer();
  5141. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  5142. int64_t cur_level = maybe_layer->layerId();
  5143. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(value, cur_level)) {
  5144. return at::_ops::fill__Tensor::call(self, value);
  5145. }
  5146. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5147. auto [value_value, value_bdim] = unwrapTensorAtLevel(value, cur_level);
  5148. batch_rule(self_value, self_bdim, value_value, value_bdim);
  5149. return self;
  5150. }
  5151. template <typename batch_rule_t, batch_rule_t batch_rule>
  5152. at::Tensor floor_generated_plumbing(const at::Tensor & self) {
  5153. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5154. auto maybe_layer = maybeCurrentDynamicLayer();
  5155. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5156. int64_t cur_level = maybe_layer->layerId();
  5157. if (!isBatchedAtLevel(self, cur_level)) {
  5158. return at::_ops::floor::call(self);
  5159. }
  5160. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5161. auto results = batch_rule(self_value, self_bdim);
  5162. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5163. }
  5164. template <typename batch_rule_t, batch_rule_t batch_rule>
  5165. at::Tensor & floor__generated_plumbing(at::Tensor & self) {
  5166. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5167. auto maybe_layer = maybeCurrentDynamicLayer();
  5168. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  5169. int64_t cur_level = maybe_layer->layerId();
  5170. if (!isBatchedAtLevel(self, cur_level)) {
  5171. return at::_ops::floor_::call(self);
  5172. }
  5173. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5174. batch_rule(self_value, self_bdim);
  5175. return self;
  5176. }
  5177. template <typename batch_rule_t, batch_rule_t batch_rule>
  5178. at::Tensor floor_divide_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  5179. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5180. auto maybe_layer = maybeCurrentDynamicLayer();
  5181. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5182. int64_t cur_level = maybe_layer->layerId();
  5183. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  5184. return at::_ops::floor_divide::call(self, other);
  5185. }
  5186. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5187. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  5188. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  5189. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5190. }
  5191. template <typename batch_rule_t, batch_rule_t batch_rule>
  5192. at::Tensor & floor_divide__Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  5193. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5194. auto maybe_layer = maybeCurrentDynamicLayer();
  5195. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  5196. int64_t cur_level = maybe_layer->layerId();
  5197. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  5198. return at::_ops::floor_divide__Tensor::call(self, other);
  5199. }
  5200. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5201. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  5202. batch_rule(self_value, self_bdim, other_value, other_bdim);
  5203. return self;
  5204. }
  5205. template <typename batch_rule_t, batch_rule_t batch_rule>
  5206. at::Tensor floor_divide_Scalar_generated_plumbing(const at::Tensor & self, const at::Scalar & other) {
  5207. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5208. auto maybe_layer = maybeCurrentDynamicLayer();
  5209. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5210. int64_t cur_level = maybe_layer->layerId();
  5211. if (!isBatchedAtLevel(self, cur_level)) {
  5212. return at::_ops::floor_divide_Scalar::call(self, other);
  5213. }
  5214. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5215. auto results = batch_rule(self_value, self_bdim, other);
  5216. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5217. }
  5218. template <typename batch_rule_t, batch_rule_t batch_rule>
  5219. at::Tensor & floor_divide__Scalar_generated_plumbing(at::Tensor & self, const at::Scalar & other) {
  5220. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5221. auto maybe_layer = maybeCurrentDynamicLayer();
  5222. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  5223. int64_t cur_level = maybe_layer->layerId();
  5224. if (!isBatchedAtLevel(self, cur_level)) {
  5225. return at::_ops::floor_divide__Scalar::call(self, other);
  5226. }
  5227. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5228. batch_rule(self_value, self_bdim, other);
  5229. return self;
  5230. }
  5231. template <typename batch_rule_t, batch_rule_t batch_rule>
  5232. at::Tensor frac_generated_plumbing(const at::Tensor & self) {
  5233. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5234. auto maybe_layer = maybeCurrentDynamicLayer();
  5235. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5236. int64_t cur_level = maybe_layer->layerId();
  5237. if (!isBatchedAtLevel(self, cur_level)) {
  5238. return at::_ops::frac::call(self);
  5239. }
  5240. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5241. auto results = batch_rule(self_value, self_bdim);
  5242. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5243. }
  5244. template <typename batch_rule_t, batch_rule_t batch_rule>
  5245. at::Tensor & frac__generated_plumbing(at::Tensor & self) {
  5246. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5247. auto maybe_layer = maybeCurrentDynamicLayer();
  5248. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  5249. int64_t cur_level = maybe_layer->layerId();
  5250. if (!isBatchedAtLevel(self, cur_level)) {
  5251. return at::_ops::frac_::call(self);
  5252. }
  5253. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5254. batch_rule(self_value, self_bdim);
  5255. return self;
  5256. }
  5257. template <typename batch_rule_t, batch_rule_t batch_rule>
  5258. at::Tensor full_like_generated_plumbing(const at::Tensor & self, const at::Scalar & fill_value, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory, ::std::optional<at::MemoryFormat> memory_format) {
  5259. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5260. auto maybe_layer = maybeCurrentDynamicLayer();
  5261. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5262. int64_t cur_level = maybe_layer->layerId();
  5263. if (!isBatchedAtLevel(self, cur_level)) {
  5264. return at::_ops::full_like::call(self, fill_value, dtype, layout, device, pin_memory, memory_format);
  5265. }
  5266. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5267. auto results = batch_rule(self_value, self_bdim, fill_value, dtype, layout, device, pin_memory, memory_format);
  5268. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5269. }
  5270. template <typename batch_rule_t, batch_rule_t batch_rule>
  5271. at::Tensor gcd_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  5272. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5273. auto maybe_layer = maybeCurrentDynamicLayer();
  5274. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5275. int64_t cur_level = maybe_layer->layerId();
  5276. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  5277. return at::_ops::gcd::call(self, other);
  5278. }
  5279. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5280. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  5281. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  5282. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5283. }
  5284. template <typename batch_rule_t, batch_rule_t batch_rule>
  5285. at::Tensor & gcd__generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  5286. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5287. auto maybe_layer = maybeCurrentDynamicLayer();
  5288. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  5289. int64_t cur_level = maybe_layer->layerId();
  5290. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  5291. return at::_ops::gcd_::call(self, other);
  5292. }
  5293. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5294. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  5295. batch_rule(self_value, self_bdim, other_value, other_bdim);
  5296. return self;
  5297. }
  5298. template <typename batch_rule_t, batch_rule_t batch_rule>
  5299. at::Tensor lcm_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  5300. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5301. auto maybe_layer = maybeCurrentDynamicLayer();
  5302. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5303. int64_t cur_level = maybe_layer->layerId();
  5304. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  5305. return at::_ops::lcm::call(self, other);
  5306. }
  5307. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5308. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  5309. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  5310. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5311. }
  5312. template <typename batch_rule_t, batch_rule_t batch_rule>
  5313. at::Tensor & lcm__generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  5314. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5315. auto maybe_layer = maybeCurrentDynamicLayer();
  5316. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  5317. int64_t cur_level = maybe_layer->layerId();
  5318. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  5319. return at::_ops::lcm_::call(self, other);
  5320. }
  5321. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5322. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  5323. batch_rule(self_value, self_bdim, other_value, other_bdim);
  5324. return self;
  5325. }
  5326. template <typename batch_rule_t, batch_rule_t batch_rule>
  5327. at::Tensor grid_sampler_generated_plumbing(const at::Tensor & input, const at::Tensor & grid, int64_t interpolation_mode, int64_t padding_mode, bool align_corners) {
  5328. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5329. auto maybe_layer = maybeCurrentDynamicLayer();
  5330. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5331. int64_t cur_level = maybe_layer->layerId();
  5332. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(grid, cur_level)) {
  5333. return at::_ops::grid_sampler::call(input, grid, interpolation_mode, padding_mode, align_corners);
  5334. }
  5335. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  5336. auto [grid_value, grid_bdim] = unwrapTensorAtLevel(grid, cur_level);
  5337. auto results = batch_rule(input_value, input_bdim, grid_value, grid_bdim, interpolation_mode, padding_mode, align_corners);
  5338. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5339. }
  5340. template <typename batch_rule_t, batch_rule_t batch_rule>
  5341. at::Tensor grid_sampler_2d_generated_plumbing(const at::Tensor & input, const at::Tensor & grid, int64_t interpolation_mode, int64_t padding_mode, bool align_corners) {
  5342. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5343. auto maybe_layer = maybeCurrentDynamicLayer();
  5344. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5345. int64_t cur_level = maybe_layer->layerId();
  5346. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(grid, cur_level)) {
  5347. return at::_ops::grid_sampler_2d::call(input, grid, interpolation_mode, padding_mode, align_corners);
  5348. }
  5349. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  5350. auto [grid_value, grid_bdim] = unwrapTensorAtLevel(grid, cur_level);
  5351. auto results = batch_rule(input_value, input_bdim, grid_value, grid_bdim, interpolation_mode, padding_mode, align_corners);
  5352. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5353. }
  5354. template <typename batch_rule_t, batch_rule_t batch_rule>
  5355. ::std::tuple<at::Tensor,at::Tensor> grid_sampler_2d_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & input, const at::Tensor & grid, int64_t interpolation_mode, int64_t padding_mode, bool align_corners, ::std::array<bool,2> output_mask) {
  5356. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5357. auto maybe_layer = maybeCurrentDynamicLayer();
  5358. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5359. int64_t cur_level = maybe_layer->layerId();
  5360. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(grid, cur_level)) {
  5361. return at::_ops::grid_sampler_2d_backward::call(grad_output, input, grid, interpolation_mode, padding_mode, align_corners, output_mask);
  5362. }
  5363. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  5364. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  5365. auto [grid_value, grid_bdim] = unwrapTensorAtLevel(grid, cur_level);
  5366. auto results = batch_rule(grad_output_value, grad_output_bdim, input_value, input_bdim, grid_value, grid_bdim, interpolation_mode, padding_mode, align_corners, output_mask);
  5367. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  5368. }
  5369. template <typename batch_rule_t, batch_rule_t batch_rule>
  5370. at::Tensor _grid_sampler_2d_cpu_fallback_generated_plumbing(const at::Tensor & input, const at::Tensor & grid, int64_t interpolation_mode, int64_t padding_mode, bool align_corners) {
  5371. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5372. auto maybe_layer = maybeCurrentDynamicLayer();
  5373. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5374. int64_t cur_level = maybe_layer->layerId();
  5375. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(grid, cur_level)) {
  5376. return at::_ops::_grid_sampler_2d_cpu_fallback::call(input, grid, interpolation_mode, padding_mode, align_corners);
  5377. }
  5378. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  5379. auto [grid_value, grid_bdim] = unwrapTensorAtLevel(grid, cur_level);
  5380. auto results = batch_rule(input_value, input_bdim, grid_value, grid_bdim, interpolation_mode, padding_mode, align_corners);
  5381. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5382. }
  5383. template <typename batch_rule_t, batch_rule_t batch_rule>
  5384. ::std::tuple<at::Tensor,at::Tensor> _grid_sampler_2d_cpu_fallback_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & input, const at::Tensor & grid, int64_t interpolation_mode, int64_t padding_mode, bool align_corners) {
  5385. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5386. auto maybe_layer = maybeCurrentDynamicLayer();
  5387. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5388. int64_t cur_level = maybe_layer->layerId();
  5389. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(grid, cur_level)) {
  5390. return at::_ops::_grid_sampler_2d_cpu_fallback_backward::call(grad_output, input, grid, interpolation_mode, padding_mode, align_corners);
  5391. }
  5392. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  5393. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  5394. auto [grid_value, grid_bdim] = unwrapTensorAtLevel(grid, cur_level);
  5395. auto results = batch_rule(grad_output_value, grad_output_bdim, input_value, input_bdim, grid_value, grid_bdim, interpolation_mode, padding_mode, align_corners);
  5396. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  5397. }
  5398. template <typename batch_rule_t, batch_rule_t batch_rule>
  5399. at::Tensor grid_sampler_3d_generated_plumbing(const at::Tensor & input, const at::Tensor & grid, int64_t interpolation_mode, int64_t padding_mode, bool align_corners) {
  5400. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5401. auto maybe_layer = maybeCurrentDynamicLayer();
  5402. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5403. int64_t cur_level = maybe_layer->layerId();
  5404. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(grid, cur_level)) {
  5405. return at::_ops::grid_sampler_3d::call(input, grid, interpolation_mode, padding_mode, align_corners);
  5406. }
  5407. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  5408. auto [grid_value, grid_bdim] = unwrapTensorAtLevel(grid, cur_level);
  5409. auto results = batch_rule(input_value, input_bdim, grid_value, grid_bdim, interpolation_mode, padding_mode, align_corners);
  5410. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5411. }
  5412. template <typename batch_rule_t, batch_rule_t batch_rule>
  5413. ::std::tuple<at::Tensor,at::Tensor> grid_sampler_3d_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & input, const at::Tensor & grid, int64_t interpolation_mode, int64_t padding_mode, bool align_corners, ::std::array<bool,2> output_mask) {
  5414. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5415. auto maybe_layer = maybeCurrentDynamicLayer();
  5416. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5417. int64_t cur_level = maybe_layer->layerId();
  5418. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(grid, cur_level)) {
  5419. return at::_ops::grid_sampler_3d_backward::call(grad_output, input, grid, interpolation_mode, padding_mode, align_corners, output_mask);
  5420. }
  5421. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  5422. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  5423. auto [grid_value, grid_bdim] = unwrapTensorAtLevel(grid, cur_level);
  5424. auto results = batch_rule(grad_output_value, grad_output_bdim, input_value, input_bdim, grid_value, grid_bdim, interpolation_mode, padding_mode, align_corners, output_mask);
  5425. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  5426. }
  5427. template <typename batch_rule_t, batch_rule_t batch_rule>
  5428. at::Tensor hinge_embedding_loss_generated_plumbing(const at::Tensor & self, const at::Tensor & target, double margin, int64_t reduction) {
  5429. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5430. auto maybe_layer = maybeCurrentDynamicLayer();
  5431. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5432. int64_t cur_level = maybe_layer->layerId();
  5433. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(target, cur_level)) {
  5434. return at::_ops::hinge_embedding_loss::call(self, target, margin, reduction);
  5435. }
  5436. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5437. auto [target_value, target_bdim] = unwrapTensorAtLevel(target, cur_level);
  5438. auto results = batch_rule(self_value, self_bdim, target_value, target_bdim, margin, reduction);
  5439. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5440. }
  5441. template <typename batch_rule_t, batch_rule_t batch_rule>
  5442. at::Tensor group_norm_generated_plumbing(const at::Tensor & input, int64_t num_groups, const ::std::optional<at::Tensor> & weight, const ::std::optional<at::Tensor> & bias, double eps, bool cudnn_enabled) {
  5443. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5444. auto maybe_layer = maybeCurrentDynamicLayer();
  5445. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5446. int64_t cur_level = maybe_layer->layerId();
  5447. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  5448. return at::_ops::group_norm::call(input, num_groups, weight, bias, eps, cudnn_enabled);
  5449. }
  5450. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  5451. std::optional<Tensor> weight_value;
  5452. std::optional<int64_t> weight_bdim;
  5453. if (weight) {
  5454. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  5455. }
  5456. std::optional<Tensor> bias_value;
  5457. std::optional<int64_t> bias_bdim;
  5458. if (bias) {
  5459. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  5460. }
  5461. auto results = batch_rule(input_value, input_bdim, num_groups, weight_value, weight_bdim, bias_value, bias_bdim, eps, cudnn_enabled);
  5462. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5463. }
  5464. template <typename batch_rule_t, batch_rule_t batch_rule>
  5465. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> native_group_norm_generated_plumbing(const at::Tensor & input, const ::std::optional<at::Tensor> & weight, const ::std::optional<at::Tensor> & bias, c10::SymInt N, c10::SymInt C, c10::SymInt HxW, int64_t group, double eps) {
  5466. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5467. auto maybe_layer = maybeCurrentDynamicLayer();
  5468. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5469. int64_t cur_level = maybe_layer->layerId();
  5470. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  5471. return at::_ops::native_group_norm::call(input, weight, bias, N, C, HxW, group, eps);
  5472. }
  5473. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  5474. std::optional<Tensor> weight_value;
  5475. std::optional<int64_t> weight_bdim;
  5476. if (weight) {
  5477. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  5478. }
  5479. std::optional<Tensor> bias_value;
  5480. std::optional<int64_t> bias_bdim;
  5481. if (bias) {
  5482. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  5483. }
  5484. auto results = batch_rule(input_value, input_bdim, weight_value, weight_bdim, bias_value, bias_bdim, N, C, HxW, group, eps);
  5485. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  5486. }
  5487. template <typename batch_rule_t, batch_rule_t batch_rule>
  5488. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> native_group_norm_backward_generated_plumbing(const at::Tensor & grad_out, const at::Tensor & input, const at::Tensor & mean, const at::Tensor & rstd, const ::std::optional<at::Tensor> & weight, c10::SymInt N, c10::SymInt C, c10::SymInt HxW, int64_t group, ::std::array<bool,3> output_mask) {
  5489. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5490. auto maybe_layer = maybeCurrentDynamicLayer();
  5491. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5492. int64_t cur_level = maybe_layer->layerId();
  5493. if (!isBatchedAtLevel(grad_out, cur_level) && !isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(mean, cur_level) && !isBatchedAtLevel(rstd, cur_level) && !isBatchedAtLevel(weight, cur_level)) {
  5494. return at::_ops::native_group_norm_backward::call(grad_out, input, mean, rstd, weight, N, C, HxW, group, output_mask);
  5495. }
  5496. auto [grad_out_value, grad_out_bdim] = unwrapTensorAtLevel(grad_out, cur_level);
  5497. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  5498. auto [mean_value, mean_bdim] = unwrapTensorAtLevel(mean, cur_level);
  5499. auto [rstd_value, rstd_bdim] = unwrapTensorAtLevel(rstd, cur_level);
  5500. std::optional<Tensor> weight_value;
  5501. std::optional<int64_t> weight_bdim;
  5502. if (weight) {
  5503. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  5504. }
  5505. auto results = batch_rule(grad_out_value, grad_out_bdim, input_value, input_bdim, mean_value, mean_bdim, rstd_value, rstd_bdim, weight_value, weight_bdim, N, C, HxW, group, output_mask);
  5506. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  5507. }
  5508. template <typename batch_rule_t, batch_rule_t batch_rule>
  5509. at::Tensor _fft_r2c_generated_plumbing(const at::Tensor & self, at::IntArrayRef dim, int64_t normalization, bool onesided) {
  5510. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5511. auto maybe_layer = maybeCurrentDynamicLayer();
  5512. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5513. int64_t cur_level = maybe_layer->layerId();
  5514. if (!isBatchedAtLevel(self, cur_level)) {
  5515. return at::_ops::_fft_r2c::call(self, dim, normalization, onesided);
  5516. }
  5517. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5518. auto results = batch_rule(self_value, self_bdim, dim, normalization, onesided);
  5519. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5520. }
  5521. template <typename batch_rule_t, batch_rule_t batch_rule>
  5522. at::Tensor _fft_c2r_generated_plumbing(const at::Tensor & self, at::IntArrayRef dim, int64_t normalization, c10::SymInt last_dim_size) {
  5523. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5524. auto maybe_layer = maybeCurrentDynamicLayer();
  5525. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5526. int64_t cur_level = maybe_layer->layerId();
  5527. if (!isBatchedAtLevel(self, cur_level)) {
  5528. return at::_ops::_fft_c2r::call(self, dim, normalization, last_dim_size);
  5529. }
  5530. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5531. auto results = batch_rule(self_value, self_bdim, dim, normalization, last_dim_size);
  5532. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5533. }
  5534. template <typename batch_rule_t, batch_rule_t batch_rule>
  5535. at::Tensor _fft_c2c_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef dim, int64_t normalization, bool forward) {
  5536. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5537. auto maybe_layer = maybeCurrentDynamicLayer();
  5538. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5539. int64_t cur_level = maybe_layer->layerId();
  5540. if (!isBatchedAtLevel(self, cur_level)) {
  5541. return at::_ops::_fft_c2c::call(self, dim, normalization, forward);
  5542. }
  5543. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5544. auto results = batch_rule(self_value, self_bdim, dim, normalization, forward);
  5545. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5546. }
  5547. template <typename batch_rule_t, batch_rule_t batch_rule>
  5548. void _validate_compressed_sparse_indices_generated_plumbing(bool is_crow, const at::Tensor & compressed_idx, const at::Tensor & plain_idx, int64_t cdim, int64_t dim, int64_t nnz) {
  5549. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5550. auto maybe_layer = maybeCurrentDynamicLayer();
  5551. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  5552. int64_t cur_level = maybe_layer->layerId();
  5553. if (!isBatchedAtLevel(compressed_idx, cur_level) && !isBatchedAtLevel(plain_idx, cur_level)) {
  5554. return at::_ops::_validate_compressed_sparse_indices::call(is_crow, compressed_idx, plain_idx, cdim, dim, nnz);
  5555. }
  5556. auto [compressed_idx_value, compressed_idx_bdim] = unwrapTensorAtLevel(compressed_idx, cur_level);
  5557. auto [plain_idx_value, plain_idx_bdim] = unwrapTensorAtLevel(plain_idx, cur_level);
  5558. batch_rule(is_crow, compressed_idx_value, compressed_idx_bdim, plain_idx_value, plain_idx_bdim, cdim, dim, nnz);
  5559. }
  5560. template <typename batch_rule_t, batch_rule_t batch_rule>
  5561. at::Tensor index_Tensor_generated_plumbing(const at::Tensor & self, const c10::List<::std::optional<at::Tensor>> & indices) {
  5562. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5563. auto maybe_layer = maybeCurrentDynamicLayer();
  5564. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5565. int64_t cur_level = maybe_layer->layerId();
  5566. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(indices, cur_level)) {
  5567. return at::_ops::index_Tensor::call(self, indices);
  5568. }
  5569. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5570. auto results = batch_rule(self_value, self_bdim, indices);
  5571. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5572. }
  5573. template <typename batch_rule_t, batch_rule_t batch_rule>
  5574. at::Tensor _unsafe_index_Tensor_generated_plumbing(const at::Tensor & self, const c10::List<::std::optional<at::Tensor>> & indices) {
  5575. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5576. auto maybe_layer = maybeCurrentDynamicLayer();
  5577. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5578. int64_t cur_level = maybe_layer->layerId();
  5579. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(indices, cur_level)) {
  5580. return at::_ops::_unsafe_index_Tensor::call(self, indices);
  5581. }
  5582. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5583. auto results = batch_rule(self_value, self_bdim, indices);
  5584. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5585. }
  5586. template <typename batch_rule_t, batch_rule_t batch_rule>
  5587. at::Tensor _unsafe_masked_index_generated_plumbing(const at::Tensor & self, const at::Tensor & mask, const c10::List<::std::optional<at::Tensor>> & indices, const at::Scalar & fill) {
  5588. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5589. auto maybe_layer = maybeCurrentDynamicLayer();
  5590. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5591. int64_t cur_level = maybe_layer->layerId();
  5592. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(mask, cur_level) && !isBatchedAtLevel(indices, cur_level)) {
  5593. return at::_ops::_unsafe_masked_index::call(self, mask, indices, fill);
  5594. }
  5595. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5596. auto [mask_value, mask_bdim] = unwrapTensorAtLevel(mask, cur_level);
  5597. auto results = batch_rule(self_value, self_bdim, mask_value, mask_bdim, indices, fill);
  5598. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5599. }
  5600. template <typename batch_rule_t, batch_rule_t batch_rule>
  5601. at::Tensor _unsafe_masked_index_put_accumulate_generated_plumbing(const at::Tensor & self, const at::Tensor & mask, const c10::List<::std::optional<at::Tensor>> & indices, const at::Tensor & values) {
  5602. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5603. auto maybe_layer = maybeCurrentDynamicLayer();
  5604. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5605. int64_t cur_level = maybe_layer->layerId();
  5606. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(mask, cur_level) && !isBatchedAtLevel(indices, cur_level) && !isBatchedAtLevel(values, cur_level)) {
  5607. return at::_ops::_unsafe_masked_index_put_accumulate::call(self, mask, indices, values);
  5608. }
  5609. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5610. auto [mask_value, mask_bdim] = unwrapTensorAtLevel(mask, cur_level);
  5611. auto [values_value, values_bdim] = unwrapTensorAtLevel(values, cur_level);
  5612. auto results = batch_rule(self_value, self_bdim, mask_value, mask_bdim, indices, values_value, values_bdim);
  5613. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5614. }
  5615. template <typename batch_rule_t, batch_rule_t batch_rule>
  5616. at::Tensor & index_copy__generated_plumbing(at::Tensor & self, int64_t dim, const at::Tensor & index, const at::Tensor & source) {
  5617. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5618. auto maybe_layer = maybeCurrentDynamicLayer();
  5619. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  5620. int64_t cur_level = maybe_layer->layerId();
  5621. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(index, cur_level) && !isBatchedAtLevel(source, cur_level)) {
  5622. return at::_ops::index_copy_::call(self, dim, index, source);
  5623. }
  5624. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5625. auto [index_value, index_bdim] = unwrapTensorAtLevel(index, cur_level);
  5626. auto [source_value, source_bdim] = unwrapTensorAtLevel(source, cur_level);
  5627. batch_rule(self_value, self_bdim, dim, index_value, index_bdim, source_value, source_bdim);
  5628. return self;
  5629. }
  5630. template <typename batch_rule_t, batch_rule_t batch_rule>
  5631. at::Tensor index_copy_generated_plumbing(const at::Tensor & self, int64_t dim, const at::Tensor & index, const at::Tensor & source) {
  5632. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5633. auto maybe_layer = maybeCurrentDynamicLayer();
  5634. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5635. int64_t cur_level = maybe_layer->layerId();
  5636. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(index, cur_level) && !isBatchedAtLevel(source, cur_level)) {
  5637. return at::_ops::index_copy::call(self, dim, index, source);
  5638. }
  5639. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5640. auto [index_value, index_bdim] = unwrapTensorAtLevel(index, cur_level);
  5641. auto [source_value, source_bdim] = unwrapTensorAtLevel(source, cur_level);
  5642. auto results = batch_rule(self_value, self_bdim, dim, index_value, index_bdim, source_value, source_bdim);
  5643. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5644. }
  5645. template <typename batch_rule_t, batch_rule_t batch_rule>
  5646. at::Tensor & index_copy__dimname_generated_plumbing(at::Tensor & self, at::Dimname dim, const at::Tensor & index, const at::Tensor & source) {
  5647. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5648. auto maybe_layer = maybeCurrentDynamicLayer();
  5649. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  5650. int64_t cur_level = maybe_layer->layerId();
  5651. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(index, cur_level) && !isBatchedAtLevel(source, cur_level)) {
  5652. return at::_ops::index_copy__dimname::call(self, dim, index, source);
  5653. }
  5654. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5655. auto [index_value, index_bdim] = unwrapTensorAtLevel(index, cur_level);
  5656. auto [source_value, source_bdim] = unwrapTensorAtLevel(source, cur_level);
  5657. batch_rule(self_value, self_bdim, dim, index_value, index_bdim, source_value, source_bdim);
  5658. return self;
  5659. }
  5660. template <typename batch_rule_t, batch_rule_t batch_rule>
  5661. at::Tensor index_copy_dimname_generated_plumbing(const at::Tensor & self, at::Dimname dim, const at::Tensor & index, const at::Tensor & source) {
  5662. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5663. auto maybe_layer = maybeCurrentDynamicLayer();
  5664. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5665. int64_t cur_level = maybe_layer->layerId();
  5666. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(index, cur_level) && !isBatchedAtLevel(source, cur_level)) {
  5667. return at::_ops::index_copy_dimname::call(self, dim, index, source);
  5668. }
  5669. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5670. auto [index_value, index_bdim] = unwrapTensorAtLevel(index, cur_level);
  5671. auto [source_value, source_bdim] = unwrapTensorAtLevel(source, cur_level);
  5672. auto results = batch_rule(self_value, self_bdim, dim, index_value, index_bdim, source_value, source_bdim);
  5673. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5674. }
  5675. template <typename batch_rule_t, batch_rule_t batch_rule>
  5676. at::Tensor & index_put__generated_plumbing(at::Tensor & self, const c10::List<::std::optional<at::Tensor>> & indices, const at::Tensor & values, bool accumulate) {
  5677. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5678. auto maybe_layer = maybeCurrentDynamicLayer();
  5679. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  5680. int64_t cur_level = maybe_layer->layerId();
  5681. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(indices, cur_level) && !isBatchedAtLevel(values, cur_level)) {
  5682. return at::_ops::index_put_::call(self, indices, values, accumulate);
  5683. }
  5684. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5685. auto [values_value, values_bdim] = unwrapTensorAtLevel(values, cur_level);
  5686. batch_rule(self_value, self_bdim, indices, values_value, values_bdim, accumulate);
  5687. return self;
  5688. }
  5689. template <typename batch_rule_t, batch_rule_t batch_rule>
  5690. at::Tensor index_put_generated_plumbing(const at::Tensor & self, const c10::List<::std::optional<at::Tensor>> & indices, const at::Tensor & values, bool accumulate) {
  5691. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5692. auto maybe_layer = maybeCurrentDynamicLayer();
  5693. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5694. int64_t cur_level = maybe_layer->layerId();
  5695. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(indices, cur_level) && !isBatchedAtLevel(values, cur_level)) {
  5696. return at::_ops::index_put::call(self, indices, values, accumulate);
  5697. }
  5698. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5699. auto [values_value, values_bdim] = unwrapTensorAtLevel(values, cur_level);
  5700. auto results = batch_rule(self_value, self_bdim, indices, values_value, values_bdim, accumulate);
  5701. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5702. }
  5703. template <typename batch_rule_t, batch_rule_t batch_rule>
  5704. at::Tensor _unsafe_index_put_generated_plumbing(const at::Tensor & self, const c10::List<::std::optional<at::Tensor>> & indices, const at::Tensor & values, bool accumulate) {
  5705. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5706. auto maybe_layer = maybeCurrentDynamicLayer();
  5707. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5708. int64_t cur_level = maybe_layer->layerId();
  5709. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(indices, cur_level) && !isBatchedAtLevel(values, cur_level)) {
  5710. return at::_ops::_unsafe_index_put::call(self, indices, values, accumulate);
  5711. }
  5712. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5713. auto [values_value, values_bdim] = unwrapTensorAtLevel(values, cur_level);
  5714. auto results = batch_rule(self_value, self_bdim, indices, values_value, values_bdim, accumulate);
  5715. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5716. }
  5717. template <typename batch_rule_t, batch_rule_t batch_rule>
  5718. at::Tensor & _index_put_impl__generated_plumbing(at::Tensor & self, const c10::List<::std::optional<at::Tensor>> & indices, const at::Tensor & values, bool accumulate, bool unsafe) {
  5719. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5720. auto maybe_layer = maybeCurrentDynamicLayer();
  5721. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  5722. int64_t cur_level = maybe_layer->layerId();
  5723. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(indices, cur_level) && !isBatchedAtLevel(values, cur_level)) {
  5724. return at::_ops::_index_put_impl_::call(self, indices, values, accumulate, unsafe);
  5725. }
  5726. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5727. auto [values_value, values_bdim] = unwrapTensorAtLevel(values, cur_level);
  5728. batch_rule(self_value, self_bdim, indices, values_value, values_bdim, accumulate, unsafe);
  5729. return self;
  5730. }
  5731. template <typename batch_rule_t, batch_rule_t batch_rule>
  5732. at::Tensor instance_norm_generated_plumbing(const at::Tensor & input, const ::std::optional<at::Tensor> & weight, const ::std::optional<at::Tensor> & bias, const ::std::optional<at::Tensor> & running_mean, const ::std::optional<at::Tensor> & running_var, bool use_input_stats, double momentum, double eps, bool cudnn_enabled) {
  5733. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5734. auto maybe_layer = maybeCurrentDynamicLayer();
  5735. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5736. int64_t cur_level = maybe_layer->layerId();
  5737. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level) && !isBatchedAtLevel(running_mean, cur_level) && !isBatchedAtLevel(running_var, cur_level)) {
  5738. return at::_ops::instance_norm::call(input, weight, bias, running_mean, running_var, use_input_stats, momentum, eps, cudnn_enabled);
  5739. }
  5740. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  5741. std::optional<Tensor> weight_value;
  5742. std::optional<int64_t> weight_bdim;
  5743. if (weight) {
  5744. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  5745. }
  5746. std::optional<Tensor> bias_value;
  5747. std::optional<int64_t> bias_bdim;
  5748. if (bias) {
  5749. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  5750. }
  5751. std::optional<Tensor> running_mean_value;
  5752. std::optional<int64_t> running_mean_bdim;
  5753. if (running_mean) {
  5754. std::tie(running_mean_value, running_mean_bdim) = unwrapTensorAtLevel(running_mean.value(), cur_level);
  5755. }
  5756. std::optional<Tensor> running_var_value;
  5757. std::optional<int64_t> running_var_bdim;
  5758. if (running_var) {
  5759. std::tie(running_var_value, running_var_bdim) = unwrapTensorAtLevel(running_var.value(), cur_level);
  5760. }
  5761. auto results = batch_rule(input_value, input_bdim, weight_value, weight_bdim, bias_value, bias_bdim, running_mean_value, running_mean_bdim, running_var_value, running_var_bdim, use_input_stats, momentum, eps, cudnn_enabled);
  5762. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5763. }
  5764. template <typename batch_rule_t, batch_rule_t batch_rule>
  5765. at::Tensor isclose_generated_plumbing(const at::Tensor & self, const at::Tensor & other, double rtol, double atol, bool equal_nan) {
  5766. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5767. auto maybe_layer = maybeCurrentDynamicLayer();
  5768. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5769. int64_t cur_level = maybe_layer->layerId();
  5770. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  5771. return at::_ops::isclose::call(self, other, rtol, atol, equal_nan);
  5772. }
  5773. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5774. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  5775. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim, rtol, atol, equal_nan);
  5776. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5777. }
  5778. template <typename batch_rule_t, batch_rule_t batch_rule>
  5779. at::Tensor isin_Tensor_Tensor_generated_plumbing(const at::Tensor & elements, const at::Tensor & test_elements, bool assume_unique, bool invert) {
  5780. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5781. auto maybe_layer = maybeCurrentDynamicLayer();
  5782. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5783. int64_t cur_level = maybe_layer->layerId();
  5784. if (!isBatchedAtLevel(elements, cur_level) && !isBatchedAtLevel(test_elements, cur_level)) {
  5785. return at::_ops::isin_Tensor_Tensor::call(elements, test_elements, assume_unique, invert);
  5786. }
  5787. auto [elements_value, elements_bdim] = unwrapTensorAtLevel(elements, cur_level);
  5788. auto [test_elements_value, test_elements_bdim] = unwrapTensorAtLevel(test_elements, cur_level);
  5789. auto results = batch_rule(elements_value, elements_bdim, test_elements_value, test_elements_bdim, assume_unique, invert);
  5790. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5791. }
  5792. template <typename batch_rule_t, batch_rule_t batch_rule>
  5793. at::Tensor isin_Tensor_Scalar_generated_plumbing(const at::Tensor & elements, const at::Scalar & test_element, bool assume_unique, bool invert) {
  5794. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5795. auto maybe_layer = maybeCurrentDynamicLayer();
  5796. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5797. int64_t cur_level = maybe_layer->layerId();
  5798. if (!isBatchedAtLevel(elements, cur_level)) {
  5799. return at::_ops::isin_Tensor_Scalar::call(elements, test_element, assume_unique, invert);
  5800. }
  5801. auto [elements_value, elements_bdim] = unwrapTensorAtLevel(elements, cur_level);
  5802. auto results = batch_rule(elements_value, elements_bdim, test_element, assume_unique, invert);
  5803. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5804. }
  5805. template <typename batch_rule_t, batch_rule_t batch_rule>
  5806. at::Tensor isin_Scalar_Tensor_generated_plumbing(const at::Scalar & element, const at::Tensor & test_elements, bool assume_unique, bool invert) {
  5807. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5808. auto maybe_layer = maybeCurrentDynamicLayer();
  5809. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5810. int64_t cur_level = maybe_layer->layerId();
  5811. if (!isBatchedAtLevel(test_elements, cur_level)) {
  5812. return at::_ops::isin_Scalar_Tensor::call(element, test_elements, assume_unique, invert);
  5813. }
  5814. auto [test_elements_value, test_elements_bdim] = unwrapTensorAtLevel(test_elements, cur_level);
  5815. auto results = batch_rule(element, test_elements_value, test_elements_bdim, assume_unique, invert);
  5816. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5817. }
  5818. template <typename batch_rule_t, batch_rule_t batch_rule>
  5819. at::Tensor isnan_generated_plumbing(const at::Tensor & self) {
  5820. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5821. auto maybe_layer = maybeCurrentDynamicLayer();
  5822. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5823. int64_t cur_level = maybe_layer->layerId();
  5824. if (!isBatchedAtLevel(self, cur_level)) {
  5825. return at::_ops::isnan::call(self);
  5826. }
  5827. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5828. auto results = batch_rule(self_value, self_bdim);
  5829. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5830. }
  5831. template <typename batch_rule_t, batch_rule_t batch_rule>
  5832. at::Tensor isreal_generated_plumbing(const at::Tensor & self) {
  5833. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5834. auto maybe_layer = maybeCurrentDynamicLayer();
  5835. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5836. int64_t cur_level = maybe_layer->layerId();
  5837. if (!isBatchedAtLevel(self, cur_level)) {
  5838. return at::_ops::isreal::call(self);
  5839. }
  5840. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5841. auto results = batch_rule(self_value, self_bdim);
  5842. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5843. }
  5844. template <typename batch_rule_t, batch_rule_t batch_rule>
  5845. at::Tensor kl_div_generated_plumbing(const at::Tensor & self, const at::Tensor & target, int64_t reduction, bool log_target) {
  5846. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5847. auto maybe_layer = maybeCurrentDynamicLayer();
  5848. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5849. int64_t cur_level = maybe_layer->layerId();
  5850. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(target, cur_level)) {
  5851. return at::_ops::kl_div::call(self, target, reduction, log_target);
  5852. }
  5853. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5854. auto [target_value, target_bdim] = unwrapTensorAtLevel(target, cur_level);
  5855. auto results = batch_rule(self_value, self_bdim, target_value, target_bdim, reduction, log_target);
  5856. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5857. }
  5858. template <typename batch_rule_t, batch_rule_t batch_rule>
  5859. at::Tensor kron_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  5860. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5861. auto maybe_layer = maybeCurrentDynamicLayer();
  5862. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5863. int64_t cur_level = maybe_layer->layerId();
  5864. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  5865. return at::_ops::kron::call(self, other);
  5866. }
  5867. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5868. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  5869. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  5870. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5871. }
  5872. template <typename batch_rule_t, batch_rule_t batch_rule>
  5873. ::std::tuple<at::Tensor,at::Tensor> kthvalue_generated_plumbing(const at::Tensor & self, c10::SymInt k, int64_t dim, bool keepdim) {
  5874. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5875. auto maybe_layer = maybeCurrentDynamicLayer();
  5876. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5877. int64_t cur_level = maybe_layer->layerId();
  5878. if (!isBatchedAtLevel(self, cur_level)) {
  5879. return at::_ops::kthvalue::call(self, k, dim, keepdim);
  5880. }
  5881. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5882. auto results = batch_rule(self_value, self_bdim, k, dim, keepdim);
  5883. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  5884. }
  5885. template <typename batch_rule_t, batch_rule_t batch_rule>
  5886. ::std::tuple<at::Tensor,at::Tensor> kthvalue_dimname_generated_plumbing(const at::Tensor & self, c10::SymInt k, at::Dimname dim, bool keepdim) {
  5887. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5888. auto maybe_layer = maybeCurrentDynamicLayer();
  5889. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5890. int64_t cur_level = maybe_layer->layerId();
  5891. if (!isBatchedAtLevel(self, cur_level)) {
  5892. return at::_ops::kthvalue_dimname::call(self, k, dim, keepdim);
  5893. }
  5894. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  5895. auto results = batch_rule(self_value, self_bdim, k, dim, keepdim);
  5896. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  5897. }
  5898. template <typename batch_rule_t, batch_rule_t batch_rule>
  5899. at::Tensor layer_norm_generated_plumbing(const at::Tensor & input, c10::SymIntArrayRef normalized_shape, const ::std::optional<at::Tensor> & weight, const ::std::optional<at::Tensor> & bias, double eps, bool cudnn_enable) {
  5900. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5901. auto maybe_layer = maybeCurrentDynamicLayer();
  5902. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5903. int64_t cur_level = maybe_layer->layerId();
  5904. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  5905. return at::_ops::layer_norm::call(input, normalized_shape, weight, bias, eps, cudnn_enable);
  5906. }
  5907. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  5908. std::optional<Tensor> weight_value;
  5909. std::optional<int64_t> weight_bdim;
  5910. if (weight) {
  5911. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  5912. }
  5913. std::optional<Tensor> bias_value;
  5914. std::optional<int64_t> bias_bdim;
  5915. if (bias) {
  5916. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  5917. }
  5918. auto results = batch_rule(input_value, input_bdim, normalized_shape, weight_value, weight_bdim, bias_value, bias_bdim, eps, cudnn_enable);
  5919. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5920. }
  5921. template <typename batch_rule_t, batch_rule_t batch_rule>
  5922. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> native_layer_norm_generated_plumbing(const at::Tensor & input, c10::SymIntArrayRef normalized_shape, const ::std::optional<at::Tensor> & weight, const ::std::optional<at::Tensor> & bias, double eps) {
  5923. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5924. auto maybe_layer = maybeCurrentDynamicLayer();
  5925. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5926. int64_t cur_level = maybe_layer->layerId();
  5927. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  5928. return at::_ops::native_layer_norm::call(input, normalized_shape, weight, bias, eps);
  5929. }
  5930. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  5931. std::optional<Tensor> weight_value;
  5932. std::optional<int64_t> weight_bdim;
  5933. if (weight) {
  5934. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  5935. }
  5936. std::optional<Tensor> bias_value;
  5937. std::optional<int64_t> bias_bdim;
  5938. if (bias) {
  5939. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  5940. }
  5941. auto results = batch_rule(input_value, input_bdim, normalized_shape, weight_value, weight_bdim, bias_value, bias_bdim, eps);
  5942. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  5943. }
  5944. template <typename batch_rule_t, batch_rule_t batch_rule>
  5945. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> native_layer_norm_backward_generated_plumbing(const at::Tensor & grad_out, const at::Tensor & input, c10::SymIntArrayRef normalized_shape, const at::Tensor & mean, const at::Tensor & rstd, const ::std::optional<at::Tensor> & weight, const ::std::optional<at::Tensor> & bias, ::std::array<bool,3> output_mask) {
  5946. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5947. auto maybe_layer = maybeCurrentDynamicLayer();
  5948. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5949. int64_t cur_level = maybe_layer->layerId();
  5950. if (!isBatchedAtLevel(grad_out, cur_level) && !isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(mean, cur_level) && !isBatchedAtLevel(rstd, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  5951. return at::_ops::native_layer_norm_backward::call(grad_out, input, normalized_shape, mean, rstd, weight, bias, output_mask);
  5952. }
  5953. auto [grad_out_value, grad_out_bdim] = unwrapTensorAtLevel(grad_out, cur_level);
  5954. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  5955. auto [mean_value, mean_bdim] = unwrapTensorAtLevel(mean, cur_level);
  5956. auto [rstd_value, rstd_bdim] = unwrapTensorAtLevel(rstd, cur_level);
  5957. std::optional<Tensor> weight_value;
  5958. std::optional<int64_t> weight_bdim;
  5959. if (weight) {
  5960. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  5961. }
  5962. std::optional<Tensor> bias_value;
  5963. std::optional<int64_t> bias_bdim;
  5964. if (bias) {
  5965. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  5966. }
  5967. auto results = batch_rule(grad_out_value, grad_out_bdim, input_value, input_bdim, normalized_shape, mean_value, mean_bdim, rstd_value, rstd_bdim, weight_value, weight_bdim, bias_value, bias_bdim, output_mask);
  5968. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  5969. }
  5970. template <typename batch_rule_t, batch_rule_t batch_rule>
  5971. at::Tensor rms_norm_generated_plumbing(const at::Tensor & input, c10::SymIntArrayRef normalized_shape, const ::std::optional<at::Tensor> & weight, ::std::optional<double> eps) {
  5972. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5973. auto maybe_layer = maybeCurrentDynamicLayer();
  5974. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5975. int64_t cur_level = maybe_layer->layerId();
  5976. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level)) {
  5977. return at::_ops::rms_norm::call(input, normalized_shape, weight, eps);
  5978. }
  5979. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  5980. std::optional<Tensor> weight_value;
  5981. std::optional<int64_t> weight_bdim;
  5982. if (weight) {
  5983. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  5984. }
  5985. auto results = batch_rule(input_value, input_bdim, normalized_shape, weight_value, weight_bdim, eps);
  5986. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  5987. }
  5988. template <typename batch_rule_t, batch_rule_t batch_rule>
  5989. ::std::tuple<at::Tensor,at::Tensor> _fused_rms_norm_generated_plumbing(const at::Tensor & input, at::IntArrayRef normalized_shape, const ::std::optional<at::Tensor> & weight, ::std::optional<double> eps) {
  5990. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  5991. auto maybe_layer = maybeCurrentDynamicLayer();
  5992. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  5993. int64_t cur_level = maybe_layer->layerId();
  5994. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level)) {
  5995. return at::_ops::_fused_rms_norm::call(input, normalized_shape, weight, eps);
  5996. }
  5997. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  5998. std::optional<Tensor> weight_value;
  5999. std::optional<int64_t> weight_bdim;
  6000. if (weight) {
  6001. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  6002. }
  6003. auto results = batch_rule(input_value, input_bdim, normalized_shape, weight_value, weight_bdim, eps);
  6004. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  6005. }
  6006. template <typename batch_rule_t, batch_rule_t batch_rule>
  6007. ::std::tuple<at::Tensor,at::Tensor> _fused_rms_norm_backward_generated_plumbing(const at::Tensor & grad_out, const at::Tensor & input, at::IntArrayRef normalized_shape, const at::Tensor & rstd, const ::std::optional<at::Tensor> & weight, ::std::array<bool,2> output_mask) {
  6008. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6009. auto maybe_layer = maybeCurrentDynamicLayer();
  6010. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6011. int64_t cur_level = maybe_layer->layerId();
  6012. if (!isBatchedAtLevel(grad_out, cur_level) && !isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(rstd, cur_level) && !isBatchedAtLevel(weight, cur_level)) {
  6013. return at::_ops::_fused_rms_norm_backward::call(grad_out, input, normalized_shape, rstd, weight, output_mask);
  6014. }
  6015. auto [grad_out_value, grad_out_bdim] = unwrapTensorAtLevel(grad_out, cur_level);
  6016. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  6017. auto [rstd_value, rstd_bdim] = unwrapTensorAtLevel(rstd, cur_level);
  6018. std::optional<Tensor> weight_value;
  6019. std::optional<int64_t> weight_bdim;
  6020. if (weight) {
  6021. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  6022. }
  6023. auto results = batch_rule(grad_out_value, grad_out_bdim, input_value, input_bdim, normalized_shape, rstd_value, rstd_bdim, weight_value, weight_bdim, output_mask);
  6024. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  6025. }
  6026. template <typename batch_rule_t, batch_rule_t batch_rule>
  6027. at::Tensor nan_to_num_generated_plumbing(const at::Tensor & self, ::std::optional<double> nan, ::std::optional<double> posinf, ::std::optional<double> neginf) {
  6028. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6029. auto maybe_layer = maybeCurrentDynamicLayer();
  6030. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6031. int64_t cur_level = maybe_layer->layerId();
  6032. if (!isBatchedAtLevel(self, cur_level)) {
  6033. return at::_ops::nan_to_num::call(self, nan, posinf, neginf);
  6034. }
  6035. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  6036. auto results = batch_rule(self_value, self_bdim, nan, posinf, neginf);
  6037. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6038. }
  6039. template <typename batch_rule_t, batch_rule_t batch_rule>
  6040. at::Tensor & nan_to_num__generated_plumbing(at::Tensor & self, ::std::optional<double> nan, ::std::optional<double> posinf, ::std::optional<double> neginf) {
  6041. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6042. auto maybe_layer = maybeCurrentDynamicLayer();
  6043. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  6044. int64_t cur_level = maybe_layer->layerId();
  6045. if (!isBatchedAtLevel(self, cur_level)) {
  6046. return at::_ops::nan_to_num_::call(self, nan, posinf, neginf);
  6047. }
  6048. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  6049. batch_rule(self_value, self_bdim, nan, posinf, neginf);
  6050. return self;
  6051. }
  6052. template <typename batch_rule_t, batch_rule_t batch_rule>
  6053. at::Tensor linear_generated_plumbing(const at::Tensor & input, const at::Tensor & weight, const ::std::optional<at::Tensor> & bias) {
  6054. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6055. auto maybe_layer = maybeCurrentDynamicLayer();
  6056. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6057. int64_t cur_level = maybe_layer->layerId();
  6058. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  6059. return at::_ops::linear::call(input, weight, bias);
  6060. }
  6061. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  6062. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  6063. std::optional<Tensor> bias_value;
  6064. std::optional<int64_t> bias_bdim;
  6065. if (bias) {
  6066. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  6067. }
  6068. auto results = batch_rule(input_value, input_bdim, weight_value, weight_bdim, bias_value, bias_bdim);
  6069. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6070. }
  6071. template <typename batch_rule_t, batch_rule_t batch_rule>
  6072. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> linear_backward_generated_plumbing(const at::Tensor & self, const at::Tensor & grad_output, const at::Tensor & weight, ::std::array<bool,3> output_mask) {
  6073. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6074. auto maybe_layer = maybeCurrentDynamicLayer();
  6075. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6076. int64_t cur_level = maybe_layer->layerId();
  6077. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(weight, cur_level)) {
  6078. return at::_ops::linear_backward::call(self, grad_output, weight, output_mask);
  6079. }
  6080. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  6081. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  6082. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  6083. auto results = batch_rule(self_value, self_bdim, grad_output_value, grad_output_bdim, weight_value, weight_bdim, output_mask);
  6084. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  6085. }
  6086. template <typename batch_rule_t, batch_rule_t batch_rule>
  6087. at::Tensor mkldnn_linear_generated_plumbing(const at::Tensor & self, const at::Tensor & weight, const ::std::optional<at::Tensor> & bias) {
  6088. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6089. auto maybe_layer = maybeCurrentDynamicLayer();
  6090. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6091. int64_t cur_level = maybe_layer->layerId();
  6092. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  6093. return at::_ops::mkldnn_linear::call(self, weight, bias);
  6094. }
  6095. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  6096. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  6097. std::optional<Tensor> bias_value;
  6098. std::optional<int64_t> bias_bdim;
  6099. if (bias) {
  6100. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  6101. }
  6102. auto results = batch_rule(self_value, self_bdim, weight_value, weight_bdim, bias_value, bias_bdim);
  6103. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6104. }
  6105. template <typename batch_rule_t, batch_rule_t batch_rule>
  6106. at::Tensor mkldnn_linear_backward_input_generated_plumbing(at::IntArrayRef input_size, const at::Tensor & grad_output, const at::Tensor & weight) {
  6107. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6108. auto maybe_layer = maybeCurrentDynamicLayer();
  6109. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6110. int64_t cur_level = maybe_layer->layerId();
  6111. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(weight, cur_level)) {
  6112. return at::_ops::mkldnn_linear_backward_input::call(input_size, grad_output, weight);
  6113. }
  6114. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  6115. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  6116. auto results = batch_rule(input_size, grad_output_value, grad_output_bdim, weight_value, weight_bdim);
  6117. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6118. }
  6119. template <typename batch_rule_t, batch_rule_t batch_rule>
  6120. ::std::tuple<at::Tensor,at::Tensor> mkldnn_linear_backward_weights_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & input, const at::Tensor & weight, bool bias_defined) {
  6121. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6122. auto maybe_layer = maybeCurrentDynamicLayer();
  6123. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6124. int64_t cur_level = maybe_layer->layerId();
  6125. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level)) {
  6126. return at::_ops::mkldnn_linear_backward_weights::call(grad_output, input, weight, bias_defined);
  6127. }
  6128. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  6129. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  6130. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  6131. auto results = batch_rule(grad_output_value, grad_output_bdim, input_value, input_bdim, weight_value, weight_bdim, bias_defined);
  6132. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  6133. }
  6134. template <typename batch_rule_t, batch_rule_t batch_rule>
  6135. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> mkldnn_linear_backward_generated_plumbing(const at::Tensor & self, const at::Tensor & grad_output, const at::Tensor & weight, ::std::array<bool,3> output_mask) {
  6136. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6137. auto maybe_layer = maybeCurrentDynamicLayer();
  6138. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6139. int64_t cur_level = maybe_layer->layerId();
  6140. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(weight, cur_level)) {
  6141. return at::_ops::mkldnn_linear_backward::call(self, grad_output, weight, output_mask);
  6142. }
  6143. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  6144. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  6145. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  6146. auto results = batch_rule(self_value, self_bdim, grad_output_value, grad_output_bdim, weight_value, weight_bdim, output_mask);
  6147. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  6148. }
  6149. template <typename batch_rule_t, batch_rule_t batch_rule>
  6150. at::Tensor _cslt_compress_generated_plumbing(const at::Tensor & input) {
  6151. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6152. auto maybe_layer = maybeCurrentDynamicLayer();
  6153. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6154. int64_t cur_level = maybe_layer->layerId();
  6155. if (!isBatchedAtLevel(input, cur_level)) {
  6156. return at::_ops::_cslt_compress::call(input);
  6157. }
  6158. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  6159. auto results = batch_rule(input_value, input_bdim);
  6160. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6161. }
  6162. template <typename batch_rule_t, batch_rule_t batch_rule>
  6163. at::Tensor _cslt_sparse_mm_generated_plumbing(const at::Tensor & compressed_A, const at::Tensor & dense_B, const ::std::optional<at::Tensor> & bias, const ::std::optional<at::Tensor> & alpha, ::std::optional<at::ScalarType> out_dtype, bool transpose_result, int64_t alg_id, int64_t split_k, int64_t split_k_mode) {
  6164. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6165. auto maybe_layer = maybeCurrentDynamicLayer();
  6166. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6167. int64_t cur_level = maybe_layer->layerId();
  6168. if (!isBatchedAtLevel(compressed_A, cur_level) && !isBatchedAtLevel(dense_B, cur_level) && !isBatchedAtLevel(bias, cur_level) && !isBatchedAtLevel(alpha, cur_level)) {
  6169. return at::_ops::_cslt_sparse_mm::call(compressed_A, dense_B, bias, alpha, out_dtype, transpose_result, alg_id, split_k, split_k_mode);
  6170. }
  6171. auto [compressed_A_value, compressed_A_bdim] = unwrapTensorAtLevel(compressed_A, cur_level);
  6172. auto [dense_B_value, dense_B_bdim] = unwrapTensorAtLevel(dense_B, cur_level);
  6173. std::optional<Tensor> bias_value;
  6174. std::optional<int64_t> bias_bdim;
  6175. if (bias) {
  6176. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  6177. }
  6178. std::optional<Tensor> alpha_value;
  6179. std::optional<int64_t> alpha_bdim;
  6180. if (alpha) {
  6181. std::tie(alpha_value, alpha_bdim) = unwrapTensorAtLevel(alpha.value(), cur_level);
  6182. }
  6183. auto results = batch_rule(compressed_A_value, compressed_A_bdim, dense_B_value, dense_B_bdim, bias_value, bias_bdim, alpha_value, alpha_bdim, out_dtype, transpose_result, alg_id, split_k, split_k_mode);
  6184. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6185. }
  6186. template <typename batch_rule_t, batch_rule_t batch_rule>
  6187. ::std::tuple<at::Tensor,at::Tensor,at::Tensor,at::Tensor,at::Tensor> _sparse_semi_structured_tile_generated_plumbing(const at::Tensor & input, c10::string_view algorithm, bool use_cutlass) {
  6188. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6189. auto maybe_layer = maybeCurrentDynamicLayer();
  6190. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6191. int64_t cur_level = maybe_layer->layerId();
  6192. if (!isBatchedAtLevel(input, cur_level)) {
  6193. return at::_ops::_sparse_semi_structured_tile::call(input, algorithm, use_cutlass);
  6194. }
  6195. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  6196. auto results = batch_rule(input_value, input_bdim, algorithm, use_cutlass);
  6197. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level), makeBatched(std::get<6>(results), std::get<7>(results), cur_level), makeBatched(std::get<8>(results), std::get<9>(results), cur_level));
  6198. }
  6199. template <typename batch_rule_t, batch_rule_t batch_rule>
  6200. ::std::tuple<at::Tensor,at::Tensor> _sparse_semi_structured_apply_generated_plumbing(const at::Tensor & input, const at::Tensor & thread_masks) {
  6201. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6202. auto maybe_layer = maybeCurrentDynamicLayer();
  6203. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6204. int64_t cur_level = maybe_layer->layerId();
  6205. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(thread_masks, cur_level)) {
  6206. return at::_ops::_sparse_semi_structured_apply::call(input, thread_masks);
  6207. }
  6208. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  6209. auto [thread_masks_value, thread_masks_bdim] = unwrapTensorAtLevel(thread_masks, cur_level);
  6210. auto results = batch_rule(input_value, input_bdim, thread_masks_value, thread_masks_bdim);
  6211. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  6212. }
  6213. template <typename batch_rule_t, batch_rule_t batch_rule>
  6214. at::Tensor _sparse_semi_structured_apply_dense_generated_plumbing(const at::Tensor & input, const at::Tensor & thread_masks) {
  6215. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6216. auto maybe_layer = maybeCurrentDynamicLayer();
  6217. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6218. int64_t cur_level = maybe_layer->layerId();
  6219. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(thread_masks, cur_level)) {
  6220. return at::_ops::_sparse_semi_structured_apply_dense::call(input, thread_masks);
  6221. }
  6222. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  6223. auto [thread_masks_value, thread_masks_bdim] = unwrapTensorAtLevel(thread_masks, cur_level);
  6224. auto results = batch_rule(input_value, input_bdim, thread_masks_value, thread_masks_bdim);
  6225. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6226. }
  6227. template <typename batch_rule_t, batch_rule_t batch_rule>
  6228. at::Tensor _sparse_semi_structured_linear_generated_plumbing(const at::Tensor & input, const at::Tensor & weight, const at::Tensor & meta, const ::std::optional<at::Tensor> & bias, ::std::optional<c10::string_view> activation, ::std::optional<at::ScalarType> out_dtype) {
  6229. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6230. auto maybe_layer = maybeCurrentDynamicLayer();
  6231. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6232. int64_t cur_level = maybe_layer->layerId();
  6233. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(meta, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  6234. return at::_ops::_sparse_semi_structured_linear::call(input, weight, meta, bias, activation, out_dtype);
  6235. }
  6236. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  6237. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  6238. auto [meta_value, meta_bdim] = unwrapTensorAtLevel(meta, cur_level);
  6239. std::optional<Tensor> bias_value;
  6240. std::optional<int64_t> bias_bdim;
  6241. if (bias) {
  6242. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  6243. }
  6244. auto results = batch_rule(input_value, input_bdim, weight_value, weight_bdim, meta_value, meta_bdim, bias_value, bias_bdim, activation, out_dtype);
  6245. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6246. }
  6247. template <typename batch_rule_t, batch_rule_t batch_rule>
  6248. at::Tensor _sparse_semi_structured_mm_generated_plumbing(const at::Tensor & mat1, const at::Tensor & mat1_meta, const at::Tensor & mat2, ::std::optional<at::ScalarType> out_dtype) {
  6249. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6250. auto maybe_layer = maybeCurrentDynamicLayer();
  6251. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6252. int64_t cur_level = maybe_layer->layerId();
  6253. if (!isBatchedAtLevel(mat1, cur_level) && !isBatchedAtLevel(mat1_meta, cur_level) && !isBatchedAtLevel(mat2, cur_level)) {
  6254. return at::_ops::_sparse_semi_structured_mm::call(mat1, mat1_meta, mat2, out_dtype);
  6255. }
  6256. auto [mat1_value, mat1_bdim] = unwrapTensorAtLevel(mat1, cur_level);
  6257. auto [mat1_meta_value, mat1_meta_bdim] = unwrapTensorAtLevel(mat1_meta, cur_level);
  6258. auto [mat2_value, mat2_bdim] = unwrapTensorAtLevel(mat2, cur_level);
  6259. auto results = batch_rule(mat1_value, mat1_bdim, mat1_meta_value, mat1_meta_bdim, mat2_value, mat2_bdim, out_dtype);
  6260. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6261. }
  6262. template <typename batch_rule_t, batch_rule_t batch_rule>
  6263. at::Tensor _sparse_semi_structured_addmm_generated_plumbing(const at::Tensor & input, const at::Tensor & mat1, const at::Tensor & mat1_meta, const at::Tensor & mat2, const at::Scalar & alpha, const at::Scalar & beta, ::std::optional<at::ScalarType> out_dtype) {
  6264. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6265. auto maybe_layer = maybeCurrentDynamicLayer();
  6266. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6267. int64_t cur_level = maybe_layer->layerId();
  6268. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(mat1, cur_level) && !isBatchedAtLevel(mat1_meta, cur_level) && !isBatchedAtLevel(mat2, cur_level)) {
  6269. return at::_ops::_sparse_semi_structured_addmm::call(input, mat1, mat1_meta, mat2, alpha, beta, out_dtype);
  6270. }
  6271. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  6272. auto [mat1_value, mat1_bdim] = unwrapTensorAtLevel(mat1, cur_level);
  6273. auto [mat1_meta_value, mat1_meta_bdim] = unwrapTensorAtLevel(mat1_meta, cur_level);
  6274. auto [mat2_value, mat2_bdim] = unwrapTensorAtLevel(mat2, cur_level);
  6275. auto results = batch_rule(input_value, input_bdim, mat1_value, mat1_bdim, mat1_meta_value, mat1_meta_bdim, mat2_value, mat2_bdim, alpha, beta, out_dtype);
  6276. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6277. }
  6278. template <typename batch_rule_t, batch_rule_t batch_rule>
  6279. at::Tensor _mixed_dtypes_linear_generated_plumbing(const at::Tensor & input, const at::Tensor & weight, const at::Tensor & scale, const ::std::optional<at::Tensor> & bias, ::std::optional<c10::string_view> activation) {
  6280. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6281. auto maybe_layer = maybeCurrentDynamicLayer();
  6282. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6283. int64_t cur_level = maybe_layer->layerId();
  6284. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(scale, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  6285. return at::_ops::_mixed_dtypes_linear::call(input, weight, scale, bias, activation);
  6286. }
  6287. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  6288. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  6289. auto [scale_value, scale_bdim] = unwrapTensorAtLevel(scale, cur_level);
  6290. std::optional<Tensor> bias_value;
  6291. std::optional<int64_t> bias_bdim;
  6292. if (bias) {
  6293. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  6294. }
  6295. auto results = batch_rule(input_value, input_bdim, weight_value, weight_bdim, scale_value, scale_bdim, bias_value, bias_bdim, activation);
  6296. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6297. }
  6298. template <typename batch_rule_t, batch_rule_t batch_rule>
  6299. at::Tensor fbgemm_linear_int8_weight_fp32_activation_generated_plumbing(const at::Tensor & input, const at::Tensor & weight, const at::Tensor & packed, const at::Tensor & col_offsets, const at::Scalar & weight_scale, const at::Scalar & weight_zero_point, const at::Tensor & bias) {
  6300. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6301. auto maybe_layer = maybeCurrentDynamicLayer();
  6302. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6303. int64_t cur_level = maybe_layer->layerId();
  6304. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(packed, cur_level) && !isBatchedAtLevel(col_offsets, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  6305. return at::_ops::fbgemm_linear_int8_weight_fp32_activation::call(input, weight, packed, col_offsets, weight_scale, weight_zero_point, bias);
  6306. }
  6307. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  6308. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  6309. auto [packed_value, packed_bdim] = unwrapTensorAtLevel(packed, cur_level);
  6310. auto [col_offsets_value, col_offsets_bdim] = unwrapTensorAtLevel(col_offsets, cur_level);
  6311. auto [bias_value, bias_bdim] = unwrapTensorAtLevel(bias, cur_level);
  6312. auto results = batch_rule(input_value, input_bdim, weight_value, weight_bdim, packed_value, packed_bdim, col_offsets_value, col_offsets_bdim, weight_scale, weight_zero_point, bias_value, bias_bdim);
  6313. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6314. }
  6315. template <typename batch_rule_t, batch_rule_t batch_rule>
  6316. at::Tensor fbgemm_linear_int8_weight_generated_plumbing(const at::Tensor & input, const at::Tensor & weight, const at::Tensor & packed, const at::Tensor & col_offsets, const at::Scalar & weight_scale, const at::Scalar & weight_zero_point, const at::Tensor & bias) {
  6317. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6318. auto maybe_layer = maybeCurrentDynamicLayer();
  6319. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6320. int64_t cur_level = maybe_layer->layerId();
  6321. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(packed, cur_level) && !isBatchedAtLevel(col_offsets, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  6322. return at::_ops::fbgemm_linear_int8_weight::call(input, weight, packed, col_offsets, weight_scale, weight_zero_point, bias);
  6323. }
  6324. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  6325. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  6326. auto [packed_value, packed_bdim] = unwrapTensorAtLevel(packed, cur_level);
  6327. auto [col_offsets_value, col_offsets_bdim] = unwrapTensorAtLevel(col_offsets, cur_level);
  6328. auto [bias_value, bias_bdim] = unwrapTensorAtLevel(bias, cur_level);
  6329. auto results = batch_rule(input_value, input_bdim, weight_value, weight_bdim, packed_value, packed_bdim, col_offsets_value, col_offsets_bdim, weight_scale, weight_zero_point, bias_value, bias_bdim);
  6330. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6331. }
  6332. template <typename batch_rule_t, batch_rule_t batch_rule>
  6333. at::Tensor fbgemm_pack_gemm_matrix_fp16_generated_plumbing(const at::Tensor & input) {
  6334. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6335. auto maybe_layer = maybeCurrentDynamicLayer();
  6336. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6337. int64_t cur_level = maybe_layer->layerId();
  6338. if (!isBatchedAtLevel(input, cur_level)) {
  6339. return at::_ops::fbgemm_pack_gemm_matrix_fp16::call(input);
  6340. }
  6341. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  6342. auto results = batch_rule(input_value, input_bdim);
  6343. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6344. }
  6345. template <typename batch_rule_t, batch_rule_t batch_rule>
  6346. at::Tensor _wrapped_linear_prepack_generated_plumbing(const at::Tensor & weight, const at::Tensor & weight_scale, const at::Tensor & weight_zero_point, const at::Tensor & bias) {
  6347. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6348. auto maybe_layer = maybeCurrentDynamicLayer();
  6349. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6350. int64_t cur_level = maybe_layer->layerId();
  6351. if (!isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(weight_scale, cur_level) && !isBatchedAtLevel(weight_zero_point, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  6352. return at::_ops::_wrapped_linear_prepack::call(weight, weight_scale, weight_zero_point, bias);
  6353. }
  6354. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  6355. auto [weight_scale_value, weight_scale_bdim] = unwrapTensorAtLevel(weight_scale, cur_level);
  6356. auto [weight_zero_point_value, weight_zero_point_bdim] = unwrapTensorAtLevel(weight_zero_point, cur_level);
  6357. auto [bias_value, bias_bdim] = unwrapTensorAtLevel(bias, cur_level);
  6358. auto results = batch_rule(weight_value, weight_bdim, weight_scale_value, weight_scale_bdim, weight_zero_point_value, weight_zero_point_bdim, bias_value, bias_bdim);
  6359. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6360. }
  6361. template <typename batch_rule_t, batch_rule_t batch_rule>
  6362. at::Tensor _wrapped_quantized_linear_prepacked_generated_plumbing(const at::Tensor & input, const at::Tensor & input_scale, const at::Tensor & input_zero_point, const at::Tensor & packed_weight, const at::Tensor & output_scale, const at::Tensor & output_zero_point, int64_t out_channel) {
  6363. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6364. auto maybe_layer = maybeCurrentDynamicLayer();
  6365. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6366. int64_t cur_level = maybe_layer->layerId();
  6367. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(input_scale, cur_level) && !isBatchedAtLevel(input_zero_point, cur_level) && !isBatchedAtLevel(packed_weight, cur_level) && !isBatchedAtLevel(output_scale, cur_level) && !isBatchedAtLevel(output_zero_point, cur_level)) {
  6368. return at::_ops::_wrapped_quantized_linear_prepacked::call(input, input_scale, input_zero_point, packed_weight, output_scale, output_zero_point, out_channel);
  6369. }
  6370. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  6371. auto [input_scale_value, input_scale_bdim] = unwrapTensorAtLevel(input_scale, cur_level);
  6372. auto [input_zero_point_value, input_zero_point_bdim] = unwrapTensorAtLevel(input_zero_point, cur_level);
  6373. auto [packed_weight_value, packed_weight_bdim] = unwrapTensorAtLevel(packed_weight, cur_level);
  6374. auto [output_scale_value, output_scale_bdim] = unwrapTensorAtLevel(output_scale, cur_level);
  6375. auto [output_zero_point_value, output_zero_point_bdim] = unwrapTensorAtLevel(output_zero_point, cur_level);
  6376. auto results = batch_rule(input_value, input_bdim, input_scale_value, input_scale_bdim, input_zero_point_value, input_zero_point_bdim, packed_weight_value, packed_weight_bdim, output_scale_value, output_scale_bdim, output_zero_point_value, output_zero_point_bdim, out_channel);
  6377. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6378. }
  6379. template <typename batch_rule_t, batch_rule_t batch_rule>
  6380. at::Tensor fbgemm_linear_fp16_weight_fp32_activation_generated_plumbing(const at::Tensor & input, const at::Tensor & packed_weight, const ::std::optional<at::Tensor> & bias) {
  6381. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6382. auto maybe_layer = maybeCurrentDynamicLayer();
  6383. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6384. int64_t cur_level = maybe_layer->layerId();
  6385. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(packed_weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  6386. return at::_ops::fbgemm_linear_fp16_weight_fp32_activation::call(input, packed_weight, bias);
  6387. }
  6388. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  6389. auto [packed_weight_value, packed_weight_bdim] = unwrapTensorAtLevel(packed_weight, cur_level);
  6390. std::optional<Tensor> bias_value;
  6391. std::optional<int64_t> bias_bdim;
  6392. if (bias) {
  6393. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  6394. }
  6395. auto results = batch_rule(input_value, input_bdim, packed_weight_value, packed_weight_bdim, bias_value, bias_bdim);
  6396. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6397. }
  6398. template <typename batch_rule_t, batch_rule_t batch_rule>
  6399. at::Tensor fbgemm_linear_fp16_weight_generated_plumbing(const at::Tensor & input, const at::Tensor & packed_weight, const at::Tensor & bias) {
  6400. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6401. auto maybe_layer = maybeCurrentDynamicLayer();
  6402. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6403. int64_t cur_level = maybe_layer->layerId();
  6404. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(packed_weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  6405. return at::_ops::fbgemm_linear_fp16_weight::call(input, packed_weight, bias);
  6406. }
  6407. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  6408. auto [packed_weight_value, packed_weight_bdim] = unwrapTensorAtLevel(packed_weight, cur_level);
  6409. auto [bias_value, bias_bdim] = unwrapTensorAtLevel(bias, cur_level);
  6410. auto results = batch_rule(input_value, input_bdim, packed_weight_value, packed_weight_bdim, bias_value, bias_bdim);
  6411. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6412. }
  6413. template <typename batch_rule_t, batch_rule_t batch_rule>
  6414. at::Tensor fbgemm_pack_quantized_matrix_generated_plumbing(const at::Tensor & input) {
  6415. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6416. auto maybe_layer = maybeCurrentDynamicLayer();
  6417. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6418. int64_t cur_level = maybe_layer->layerId();
  6419. if (!isBatchedAtLevel(input, cur_level)) {
  6420. return at::_ops::fbgemm_pack_quantized_matrix::call(input);
  6421. }
  6422. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  6423. auto results = batch_rule(input_value, input_bdim);
  6424. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6425. }
  6426. template <typename batch_rule_t, batch_rule_t batch_rule>
  6427. at::Tensor fbgemm_pack_quantized_matrix_KN_generated_plumbing(const at::Tensor & input, int64_t K, int64_t N) {
  6428. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6429. auto maybe_layer = maybeCurrentDynamicLayer();
  6430. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6431. int64_t cur_level = maybe_layer->layerId();
  6432. if (!isBatchedAtLevel(input, cur_level)) {
  6433. return at::_ops::fbgemm_pack_quantized_matrix_KN::call(input, K, N);
  6434. }
  6435. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  6436. auto results = batch_rule(input_value, input_bdim, K, N);
  6437. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6438. }
  6439. template <typename batch_rule_t, batch_rule_t batch_rule>
  6440. at::Tensor ldexp_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  6441. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6442. auto maybe_layer = maybeCurrentDynamicLayer();
  6443. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6444. int64_t cur_level = maybe_layer->layerId();
  6445. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  6446. return at::_ops::ldexp_Tensor::call(self, other);
  6447. }
  6448. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  6449. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  6450. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  6451. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6452. }
  6453. template <typename batch_rule_t, batch_rule_t batch_rule>
  6454. at::Tensor & ldexp__generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  6455. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6456. auto maybe_layer = maybeCurrentDynamicLayer();
  6457. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  6458. int64_t cur_level = maybe_layer->layerId();
  6459. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  6460. return at::_ops::ldexp_::call(self, other);
  6461. }
  6462. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  6463. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  6464. batch_rule(self_value, self_bdim, other_value, other_bdim);
  6465. return self;
  6466. }
  6467. template <typename batch_rule_t, batch_rule_t batch_rule>
  6468. at::Tensor linspace_Tensor_Tensor_generated_plumbing(const at::Tensor & start, const at::Tensor & end, int64_t steps, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory) {
  6469. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6470. auto maybe_layer = maybeCurrentDynamicLayer();
  6471. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6472. int64_t cur_level = maybe_layer->layerId();
  6473. if (!isBatchedAtLevel(start, cur_level) && !isBatchedAtLevel(end, cur_level)) {
  6474. return at::_ops::linspace_Tensor_Tensor::call(start, end, steps, dtype, layout, device, pin_memory);
  6475. }
  6476. auto [start_value, start_bdim] = unwrapTensorAtLevel(start, cur_level);
  6477. auto [end_value, end_bdim] = unwrapTensorAtLevel(end, cur_level);
  6478. auto results = batch_rule(start_value, start_bdim, end_value, end_bdim, steps, dtype, layout, device, pin_memory);
  6479. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6480. }
  6481. template <typename batch_rule_t, batch_rule_t batch_rule>
  6482. at::Tensor linspace_Tensor_Scalar_generated_plumbing(const at::Tensor & start, const at::Scalar & end, int64_t steps, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory) {
  6483. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6484. auto maybe_layer = maybeCurrentDynamicLayer();
  6485. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6486. int64_t cur_level = maybe_layer->layerId();
  6487. if (!isBatchedAtLevel(start, cur_level)) {
  6488. return at::_ops::linspace_Tensor_Scalar::call(start, end, steps, dtype, layout, device, pin_memory);
  6489. }
  6490. auto [start_value, start_bdim] = unwrapTensorAtLevel(start, cur_level);
  6491. auto results = batch_rule(start_value, start_bdim, end, steps, dtype, layout, device, pin_memory);
  6492. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6493. }
  6494. template <typename batch_rule_t, batch_rule_t batch_rule>
  6495. at::Tensor linspace_Scalar_Tensor_generated_plumbing(const at::Scalar & start, const at::Tensor & end, int64_t steps, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory) {
  6496. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6497. auto maybe_layer = maybeCurrentDynamicLayer();
  6498. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6499. int64_t cur_level = maybe_layer->layerId();
  6500. if (!isBatchedAtLevel(end, cur_level)) {
  6501. return at::_ops::linspace_Scalar_Tensor::call(start, end, steps, dtype, layout, device, pin_memory);
  6502. }
  6503. auto [end_value, end_bdim] = unwrapTensorAtLevel(end, cur_level);
  6504. auto results = batch_rule(start, end_value, end_bdim, steps, dtype, layout, device, pin_memory);
  6505. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6506. }
  6507. template <typename batch_rule_t, batch_rule_t batch_rule>
  6508. at::Tensor log_generated_plumbing(const at::Tensor & self) {
  6509. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6510. auto maybe_layer = maybeCurrentDynamicLayer();
  6511. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6512. int64_t cur_level = maybe_layer->layerId();
  6513. if (!isBatchedAtLevel(self, cur_level)) {
  6514. return at::_ops::log::call(self);
  6515. }
  6516. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  6517. auto results = batch_rule(self_value, self_bdim);
  6518. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6519. }
  6520. template <typename batch_rule_t, batch_rule_t batch_rule>
  6521. at::Tensor & log__generated_plumbing(at::Tensor & self) {
  6522. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6523. auto maybe_layer = maybeCurrentDynamicLayer();
  6524. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  6525. int64_t cur_level = maybe_layer->layerId();
  6526. if (!isBatchedAtLevel(self, cur_level)) {
  6527. return at::_ops::log_::call(self);
  6528. }
  6529. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  6530. batch_rule(self_value, self_bdim);
  6531. return self;
  6532. }
  6533. template <typename batch_rule_t, batch_rule_t batch_rule>
  6534. at::Tensor log10_generated_plumbing(const at::Tensor & self) {
  6535. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6536. auto maybe_layer = maybeCurrentDynamicLayer();
  6537. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6538. int64_t cur_level = maybe_layer->layerId();
  6539. if (!isBatchedAtLevel(self, cur_level)) {
  6540. return at::_ops::log10::call(self);
  6541. }
  6542. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  6543. auto results = batch_rule(self_value, self_bdim);
  6544. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6545. }
  6546. template <typename batch_rule_t, batch_rule_t batch_rule>
  6547. at::Tensor & log10__generated_plumbing(at::Tensor & self) {
  6548. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6549. auto maybe_layer = maybeCurrentDynamicLayer();
  6550. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  6551. int64_t cur_level = maybe_layer->layerId();
  6552. if (!isBatchedAtLevel(self, cur_level)) {
  6553. return at::_ops::log10_::call(self);
  6554. }
  6555. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  6556. batch_rule(self_value, self_bdim);
  6557. return self;
  6558. }
  6559. template <typename batch_rule_t, batch_rule_t batch_rule>
  6560. at::Tensor log1p_generated_plumbing(const at::Tensor & self) {
  6561. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6562. auto maybe_layer = maybeCurrentDynamicLayer();
  6563. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6564. int64_t cur_level = maybe_layer->layerId();
  6565. if (!isBatchedAtLevel(self, cur_level)) {
  6566. return at::_ops::log1p::call(self);
  6567. }
  6568. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  6569. auto results = batch_rule(self_value, self_bdim);
  6570. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6571. }
  6572. template <typename batch_rule_t, batch_rule_t batch_rule>
  6573. at::Tensor & log1p__generated_plumbing(at::Tensor & self) {
  6574. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6575. auto maybe_layer = maybeCurrentDynamicLayer();
  6576. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  6577. int64_t cur_level = maybe_layer->layerId();
  6578. if (!isBatchedAtLevel(self, cur_level)) {
  6579. return at::_ops::log1p_::call(self);
  6580. }
  6581. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  6582. batch_rule(self_value, self_bdim);
  6583. return self;
  6584. }
  6585. template <typename batch_rule_t, batch_rule_t batch_rule>
  6586. at::Tensor log2_generated_plumbing(const at::Tensor & self) {
  6587. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6588. auto maybe_layer = maybeCurrentDynamicLayer();
  6589. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6590. int64_t cur_level = maybe_layer->layerId();
  6591. if (!isBatchedAtLevel(self, cur_level)) {
  6592. return at::_ops::log2::call(self);
  6593. }
  6594. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  6595. auto results = batch_rule(self_value, self_bdim);
  6596. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6597. }
  6598. template <typename batch_rule_t, batch_rule_t batch_rule>
  6599. at::Tensor & log2__generated_plumbing(at::Tensor & self) {
  6600. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6601. auto maybe_layer = maybeCurrentDynamicLayer();
  6602. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  6603. int64_t cur_level = maybe_layer->layerId();
  6604. if (!isBatchedAtLevel(self, cur_level)) {
  6605. return at::_ops::log2_::call(self);
  6606. }
  6607. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  6608. batch_rule(self_value, self_bdim);
  6609. return self;
  6610. }
  6611. template <typename batch_rule_t, batch_rule_t batch_rule>
  6612. at::Tensor logaddexp_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  6613. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6614. auto maybe_layer = maybeCurrentDynamicLayer();
  6615. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6616. int64_t cur_level = maybe_layer->layerId();
  6617. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  6618. return at::_ops::logaddexp::call(self, other);
  6619. }
  6620. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  6621. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  6622. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  6623. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6624. }
  6625. template <typename batch_rule_t, batch_rule_t batch_rule>
  6626. at::Tensor logaddexp2_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  6627. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6628. auto maybe_layer = maybeCurrentDynamicLayer();
  6629. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6630. int64_t cur_level = maybe_layer->layerId();
  6631. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  6632. return at::_ops::logaddexp2::call(self, other);
  6633. }
  6634. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  6635. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  6636. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  6637. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6638. }
  6639. template <typename batch_rule_t, batch_rule_t batch_rule>
  6640. at::Tensor xlogy_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  6641. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6642. auto maybe_layer = maybeCurrentDynamicLayer();
  6643. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6644. int64_t cur_level = maybe_layer->layerId();
  6645. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  6646. return at::_ops::xlogy_Tensor::call(self, other);
  6647. }
  6648. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  6649. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  6650. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  6651. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6652. }
  6653. template <typename batch_rule_t, batch_rule_t batch_rule>
  6654. at::Tensor xlogy_Scalar_Self_generated_plumbing(const at::Scalar & self, const at::Tensor & other) {
  6655. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6656. auto maybe_layer = maybeCurrentDynamicLayer();
  6657. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6658. int64_t cur_level = maybe_layer->layerId();
  6659. if (!isBatchedAtLevel(other, cur_level)) {
  6660. return at::_ops::xlogy_Scalar_Self::call(self, other);
  6661. }
  6662. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  6663. auto results = batch_rule(self, other_value, other_bdim);
  6664. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6665. }
  6666. template <typename batch_rule_t, batch_rule_t batch_rule>
  6667. at::Tensor xlogy_Scalar_Other_generated_plumbing(const at::Tensor & self, const at::Scalar & other) {
  6668. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6669. auto maybe_layer = maybeCurrentDynamicLayer();
  6670. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6671. int64_t cur_level = maybe_layer->layerId();
  6672. if (!isBatchedAtLevel(self, cur_level)) {
  6673. return at::_ops::xlogy_Scalar_Other::call(self, other);
  6674. }
  6675. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  6676. auto results = batch_rule(self_value, self_bdim, other);
  6677. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6678. }
  6679. template <typename batch_rule_t, batch_rule_t batch_rule>
  6680. at::Tensor & xlogy__Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  6681. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6682. auto maybe_layer = maybeCurrentDynamicLayer();
  6683. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  6684. int64_t cur_level = maybe_layer->layerId();
  6685. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  6686. return at::_ops::xlogy__Tensor::call(self, other);
  6687. }
  6688. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  6689. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  6690. batch_rule(self_value, self_bdim, other_value, other_bdim);
  6691. return self;
  6692. }
  6693. template <typename batch_rule_t, batch_rule_t batch_rule>
  6694. at::Tensor & xlogy__Scalar_Other_generated_plumbing(at::Tensor & self, const at::Scalar & other) {
  6695. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6696. auto maybe_layer = maybeCurrentDynamicLayer();
  6697. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  6698. int64_t cur_level = maybe_layer->layerId();
  6699. if (!isBatchedAtLevel(self, cur_level)) {
  6700. return at::_ops::xlogy__Scalar_Other::call(self, other);
  6701. }
  6702. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  6703. batch_rule(self_value, self_bdim, other);
  6704. return self;
  6705. }
  6706. template <typename batch_rule_t, batch_rule_t batch_rule>
  6707. at::Tensor logspace_Tensor_Tensor_generated_plumbing(const at::Tensor & start, const at::Tensor & end, int64_t steps, double base, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory) {
  6708. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6709. auto maybe_layer = maybeCurrentDynamicLayer();
  6710. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6711. int64_t cur_level = maybe_layer->layerId();
  6712. if (!isBatchedAtLevel(start, cur_level) && !isBatchedAtLevel(end, cur_level)) {
  6713. return at::_ops::logspace_Tensor_Tensor::call(start, end, steps, base, dtype, layout, device, pin_memory);
  6714. }
  6715. auto [start_value, start_bdim] = unwrapTensorAtLevel(start, cur_level);
  6716. auto [end_value, end_bdim] = unwrapTensorAtLevel(end, cur_level);
  6717. auto results = batch_rule(start_value, start_bdim, end_value, end_bdim, steps, base, dtype, layout, device, pin_memory);
  6718. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6719. }
  6720. template <typename batch_rule_t, batch_rule_t batch_rule>
  6721. at::Tensor logspace_Tensor_Scalar_generated_plumbing(const at::Tensor & start, const at::Scalar & end, int64_t steps, double base, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory) {
  6722. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6723. auto maybe_layer = maybeCurrentDynamicLayer();
  6724. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6725. int64_t cur_level = maybe_layer->layerId();
  6726. if (!isBatchedAtLevel(start, cur_level)) {
  6727. return at::_ops::logspace_Tensor_Scalar::call(start, end, steps, base, dtype, layout, device, pin_memory);
  6728. }
  6729. auto [start_value, start_bdim] = unwrapTensorAtLevel(start, cur_level);
  6730. auto results = batch_rule(start_value, start_bdim, end, steps, base, dtype, layout, device, pin_memory);
  6731. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6732. }
  6733. template <typename batch_rule_t, batch_rule_t batch_rule>
  6734. at::Tensor logspace_Scalar_Tensor_generated_plumbing(const at::Scalar & start, const at::Tensor & end, int64_t steps, double base, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory) {
  6735. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6736. auto maybe_layer = maybeCurrentDynamicLayer();
  6737. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6738. int64_t cur_level = maybe_layer->layerId();
  6739. if (!isBatchedAtLevel(end, cur_level)) {
  6740. return at::_ops::logspace_Scalar_Tensor::call(start, end, steps, base, dtype, layout, device, pin_memory);
  6741. }
  6742. auto [end_value, end_bdim] = unwrapTensorAtLevel(end, cur_level);
  6743. auto results = batch_rule(start, end_value, end_bdim, steps, base, dtype, layout, device, pin_memory);
  6744. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6745. }
  6746. template <typename batch_rule_t, batch_rule_t batch_rule>
  6747. at::Tensor log_softmax_int_generated_plumbing(const at::Tensor & self, int64_t dim, ::std::optional<at::ScalarType> dtype) {
  6748. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6749. auto maybe_layer = maybeCurrentDynamicLayer();
  6750. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6751. int64_t cur_level = maybe_layer->layerId();
  6752. if (!isBatchedAtLevel(self, cur_level)) {
  6753. return at::_ops::log_softmax_int::call(self, dim, dtype);
  6754. }
  6755. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  6756. auto results = batch_rule(self_value, self_bdim, dim, dtype);
  6757. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6758. }
  6759. template <typename batch_rule_t, batch_rule_t batch_rule>
  6760. at::Tensor log_softmax_Dimname_generated_plumbing(const at::Tensor & self, at::Dimname dim, ::std::optional<at::ScalarType> dtype) {
  6761. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6762. auto maybe_layer = maybeCurrentDynamicLayer();
  6763. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6764. int64_t cur_level = maybe_layer->layerId();
  6765. if (!isBatchedAtLevel(self, cur_level)) {
  6766. return at::_ops::log_softmax_Dimname::call(self, dim, dtype);
  6767. }
  6768. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  6769. auto results = batch_rule(self_value, self_bdim, dim, dtype);
  6770. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6771. }
  6772. template <typename batch_rule_t, batch_rule_t batch_rule>
  6773. at::Tensor _log_softmax_generated_plumbing(const at::Tensor & self, int64_t dim, bool half_to_float) {
  6774. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6775. auto maybe_layer = maybeCurrentDynamicLayer();
  6776. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6777. int64_t cur_level = maybe_layer->layerId();
  6778. if (!isBatchedAtLevel(self, cur_level)) {
  6779. return at::_ops::_log_softmax::call(self, dim, half_to_float);
  6780. }
  6781. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  6782. auto results = batch_rule(self_value, self_bdim, dim, half_to_float);
  6783. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6784. }
  6785. template <typename batch_rule_t, batch_rule_t batch_rule>
  6786. at::Tensor _log_softmax_backward_data_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & output, int64_t dim, at::ScalarType input_dtype) {
  6787. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6788. auto maybe_layer = maybeCurrentDynamicLayer();
  6789. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6790. int64_t cur_level = maybe_layer->layerId();
  6791. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(output, cur_level)) {
  6792. return at::_ops::_log_softmax_backward_data::call(grad_output, output, dim, input_dtype);
  6793. }
  6794. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  6795. auto [output_value, output_bdim] = unwrapTensorAtLevel(output, cur_level);
  6796. auto results = batch_rule(grad_output_value, grad_output_bdim, output_value, output_bdim, dim, input_dtype);
  6797. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6798. }
  6799. template <typename batch_rule_t, batch_rule_t batch_rule>
  6800. at::Tensor _logcumsumexp_generated_plumbing(const at::Tensor & self, int64_t dim) {
  6801. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6802. auto maybe_layer = maybeCurrentDynamicLayer();
  6803. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6804. int64_t cur_level = maybe_layer->layerId();
  6805. if (!isBatchedAtLevel(self, cur_level)) {
  6806. return at::_ops::_logcumsumexp::call(self, dim);
  6807. }
  6808. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  6809. auto results = batch_rule(self_value, self_bdim, dim);
  6810. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6811. }
  6812. template <typename batch_rule_t, batch_rule_t batch_rule>
  6813. at::Tensor logcumsumexp_generated_plumbing(const at::Tensor & self, int64_t dim) {
  6814. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6815. auto maybe_layer = maybeCurrentDynamicLayer();
  6816. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6817. int64_t cur_level = maybe_layer->layerId();
  6818. if (!isBatchedAtLevel(self, cur_level)) {
  6819. return at::_ops::logcumsumexp::call(self, dim);
  6820. }
  6821. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  6822. auto results = batch_rule(self_value, self_bdim, dim);
  6823. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6824. }
  6825. template <typename batch_rule_t, batch_rule_t batch_rule>
  6826. at::Tensor logcumsumexp_dimname_generated_plumbing(const at::Tensor & self, at::Dimname dim) {
  6827. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6828. auto maybe_layer = maybeCurrentDynamicLayer();
  6829. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6830. int64_t cur_level = maybe_layer->layerId();
  6831. if (!isBatchedAtLevel(self, cur_level)) {
  6832. return at::_ops::logcumsumexp_dimname::call(self, dim);
  6833. }
  6834. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  6835. auto results = batch_rule(self_value, self_bdim, dim);
  6836. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6837. }
  6838. template <typename batch_rule_t, batch_rule_t batch_rule>
  6839. at::Tensor logsumexp_generated_plumbing(const at::Tensor & self, at::IntArrayRef dim, bool keepdim) {
  6840. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6841. auto maybe_layer = maybeCurrentDynamicLayer();
  6842. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6843. int64_t cur_level = maybe_layer->layerId();
  6844. if (!isBatchedAtLevel(self, cur_level)) {
  6845. return at::_ops::logsumexp::call(self, dim, keepdim);
  6846. }
  6847. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  6848. auto results = batch_rule(self_value, self_bdim, dim, keepdim);
  6849. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6850. }
  6851. template <typename batch_rule_t, batch_rule_t batch_rule>
  6852. at::Tensor logsumexp_names_generated_plumbing(const at::Tensor & self, at::DimnameList dim, bool keepdim) {
  6853. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6854. auto maybe_layer = maybeCurrentDynamicLayer();
  6855. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6856. int64_t cur_level = maybe_layer->layerId();
  6857. if (!isBatchedAtLevel(self, cur_level)) {
  6858. return at::_ops::logsumexp_names::call(self, dim, keepdim);
  6859. }
  6860. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  6861. auto results = batch_rule(self_value, self_bdim, dim, keepdim);
  6862. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6863. }
  6864. template <typename batch_rule_t, batch_rule_t batch_rule>
  6865. at::Tensor margin_ranking_loss_generated_plumbing(const at::Tensor & input1, const at::Tensor & input2, const at::Tensor & target, double margin, int64_t reduction) {
  6866. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6867. auto maybe_layer = maybeCurrentDynamicLayer();
  6868. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6869. int64_t cur_level = maybe_layer->layerId();
  6870. if (!isBatchedAtLevel(input1, cur_level) && !isBatchedAtLevel(input2, cur_level) && !isBatchedAtLevel(target, cur_level)) {
  6871. return at::_ops::margin_ranking_loss::call(input1, input2, target, margin, reduction);
  6872. }
  6873. auto [input1_value, input1_bdim] = unwrapTensorAtLevel(input1, cur_level);
  6874. auto [input2_value, input2_bdim] = unwrapTensorAtLevel(input2, cur_level);
  6875. auto [target_value, target_bdim] = unwrapTensorAtLevel(target, cur_level);
  6876. auto results = batch_rule(input1_value, input1_bdim, input2_value, input2_bdim, target_value, target_bdim, margin, reduction);
  6877. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6878. }
  6879. template <typename batch_rule_t, batch_rule_t batch_rule>
  6880. at::Tensor matmul_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  6881. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6882. auto maybe_layer = maybeCurrentDynamicLayer();
  6883. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6884. int64_t cur_level = maybe_layer->layerId();
  6885. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  6886. return at::_ops::matmul::call(self, other);
  6887. }
  6888. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  6889. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  6890. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  6891. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6892. }
  6893. template <typename batch_rule_t, batch_rule_t batch_rule>
  6894. ::std::tuple<at::Tensor,at::Tensor> matmul_backward_generated_plumbing(const at::Tensor & grad, const at::Tensor & self, const at::Tensor & other, ::std::array<bool,2> mask) {
  6895. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6896. auto maybe_layer = maybeCurrentDynamicLayer();
  6897. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6898. int64_t cur_level = maybe_layer->layerId();
  6899. if (!isBatchedAtLevel(grad, cur_level) && !isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  6900. return at::_ops::matmul_backward::call(grad, self, other, mask);
  6901. }
  6902. auto [grad_value, grad_bdim] = unwrapTensorAtLevel(grad, cur_level);
  6903. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  6904. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  6905. auto results = batch_rule(grad_value, grad_bdim, self_value, self_bdim, other_value, other_bdim, mask);
  6906. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  6907. }
  6908. template <typename batch_rule_t, batch_rule_t batch_rule>
  6909. at::Tensor matrix_power_generated_plumbing(const at::Tensor & self, int64_t n) {
  6910. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6911. auto maybe_layer = maybeCurrentDynamicLayer();
  6912. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6913. int64_t cur_level = maybe_layer->layerId();
  6914. if (!isBatchedAtLevel(self, cur_level)) {
  6915. return at::_ops::matrix_power::call(self, n);
  6916. }
  6917. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  6918. auto results = batch_rule(self_value, self_bdim, n);
  6919. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6920. }
  6921. template <typename batch_rule_t, batch_rule_t batch_rule>
  6922. at::Tensor matrix_exp_generated_plumbing(const at::Tensor & self) {
  6923. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6924. auto maybe_layer = maybeCurrentDynamicLayer();
  6925. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6926. int64_t cur_level = maybe_layer->layerId();
  6927. if (!isBatchedAtLevel(self, cur_level)) {
  6928. return at::_ops::matrix_exp::call(self);
  6929. }
  6930. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  6931. auto results = batch_rule(self_value, self_bdim);
  6932. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6933. }
  6934. template <typename batch_rule_t, batch_rule_t batch_rule>
  6935. at::Tensor matrix_exp_backward_generated_plumbing(const at::Tensor & self, const at::Tensor & grad) {
  6936. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6937. auto maybe_layer = maybeCurrentDynamicLayer();
  6938. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6939. int64_t cur_level = maybe_layer->layerId();
  6940. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(grad, cur_level)) {
  6941. return at::_ops::matrix_exp_backward::call(self, grad);
  6942. }
  6943. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  6944. auto [grad_value, grad_bdim] = unwrapTensorAtLevel(grad, cur_level);
  6945. auto results = batch_rule(self_value, self_bdim, grad_value, grad_bdim);
  6946. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  6947. }
  6948. template <typename batch_rule_t, batch_rule_t batch_rule>
  6949. ::std::tuple<at::Tensor,at::Tensor> _aminmax_generated_plumbing(const at::Tensor & self) {
  6950. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6951. auto maybe_layer = maybeCurrentDynamicLayer();
  6952. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6953. int64_t cur_level = maybe_layer->layerId();
  6954. if (!isBatchedAtLevel(self, cur_level)) {
  6955. return at::_ops::_aminmax::call(self);
  6956. }
  6957. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  6958. auto results = batch_rule(self_value, self_bdim);
  6959. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  6960. }
  6961. template <typename batch_rule_t, batch_rule_t batch_rule>
  6962. ::std::tuple<at::Tensor,at::Tensor> _aminmax_dim_generated_plumbing(const at::Tensor & self, int64_t dim, bool keepdim) {
  6963. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6964. auto maybe_layer = maybeCurrentDynamicLayer();
  6965. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6966. int64_t cur_level = maybe_layer->layerId();
  6967. if (!isBatchedAtLevel(self, cur_level)) {
  6968. return at::_ops::_aminmax_dim::call(self, dim, keepdim);
  6969. }
  6970. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  6971. auto results = batch_rule(self_value, self_bdim, dim, keepdim);
  6972. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  6973. }
  6974. template <typename batch_rule_t, batch_rule_t batch_rule>
  6975. ::std::tuple<at::Tensor,at::Tensor> aminmax_generated_plumbing(const at::Tensor & self, ::std::optional<int64_t> dim, bool keepdim) {
  6976. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6977. auto maybe_layer = maybeCurrentDynamicLayer();
  6978. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6979. int64_t cur_level = maybe_layer->layerId();
  6980. if (!isBatchedAtLevel(self, cur_level)) {
  6981. return at::_ops::aminmax::call(self, dim, keepdim);
  6982. }
  6983. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  6984. auto results = batch_rule(self_value, self_bdim, dim, keepdim);
  6985. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  6986. }
  6987. template <typename batch_rule_t, batch_rule_t batch_rule>
  6988. at::Tensor _compute_linear_combination_generated_plumbing(const at::Tensor & input, const at::Tensor & coefficients) {
  6989. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  6990. auto maybe_layer = maybeCurrentDynamicLayer();
  6991. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  6992. int64_t cur_level = maybe_layer->layerId();
  6993. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(coefficients, cur_level)) {
  6994. return at::_ops::_compute_linear_combination::call(input, coefficients);
  6995. }
  6996. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  6997. auto [coefficients_value, coefficients_bdim] = unwrapTensorAtLevel(coefficients, cur_level);
  6998. auto results = batch_rule(input_value, input_bdim, coefficients_value, coefficients_bdim);
  6999. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7000. }
  7001. template <typename batch_rule_t, batch_rule_t batch_rule>
  7002. ::std::tuple<at::Tensor,at::Tensor> max_dim_generated_plumbing(const at::Tensor & self, int64_t dim, bool keepdim) {
  7003. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7004. auto maybe_layer = maybeCurrentDynamicLayer();
  7005. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7006. int64_t cur_level = maybe_layer->layerId();
  7007. if (!isBatchedAtLevel(self, cur_level)) {
  7008. return at::_ops::max_dim::call(self, dim, keepdim);
  7009. }
  7010. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7011. auto results = batch_rule(self_value, self_bdim, dim, keepdim);
  7012. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  7013. }
  7014. template <typename batch_rule_t, batch_rule_t batch_rule>
  7015. ::std::tuple<at::Tensor,at::Tensor> max_names_dim_generated_plumbing(const at::Tensor & self, at::Dimname dim, bool keepdim) {
  7016. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7017. auto maybe_layer = maybeCurrentDynamicLayer();
  7018. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7019. int64_t cur_level = maybe_layer->layerId();
  7020. if (!isBatchedAtLevel(self, cur_level)) {
  7021. return at::_ops::max_names_dim::call(self, dim, keepdim);
  7022. }
  7023. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7024. auto results = batch_rule(self_value, self_bdim, dim, keepdim);
  7025. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  7026. }
  7027. template <typename batch_rule_t, batch_rule_t batch_rule>
  7028. at::Tensor value_selecting_reduction_backward_generated_plumbing(const at::Tensor & grad, int64_t dim, const at::Tensor & indices, c10::SymIntArrayRef sizes, bool keepdim) {
  7029. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7030. auto maybe_layer = maybeCurrentDynamicLayer();
  7031. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7032. int64_t cur_level = maybe_layer->layerId();
  7033. if (!isBatchedAtLevel(grad, cur_level) && !isBatchedAtLevel(indices, cur_level)) {
  7034. return at::_ops::value_selecting_reduction_backward::call(grad, dim, indices, sizes, keepdim);
  7035. }
  7036. auto [grad_value, grad_bdim] = unwrapTensorAtLevel(grad, cur_level);
  7037. auto [indices_value, indices_bdim] = unwrapTensorAtLevel(indices, cur_level);
  7038. auto results = batch_rule(grad_value, grad_bdim, dim, indices_value, indices_bdim, sizes, keepdim);
  7039. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7040. }
  7041. template <typename batch_rule_t, batch_rule_t batch_rule>
  7042. at::Tensor amax_generated_plumbing(const at::Tensor & self, at::IntArrayRef dim, bool keepdim) {
  7043. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7044. auto maybe_layer = maybeCurrentDynamicLayer();
  7045. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7046. int64_t cur_level = maybe_layer->layerId();
  7047. if (!isBatchedAtLevel(self, cur_level)) {
  7048. return at::_ops::amax::call(self, dim, keepdim);
  7049. }
  7050. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7051. auto results = batch_rule(self_value, self_bdim, dim, keepdim);
  7052. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7053. }
  7054. template <typename batch_rule_t, batch_rule_t batch_rule>
  7055. ::std::tuple<at::Tensor,at::Tensor> max_pool1d_with_indices_generated_plumbing(const at::Tensor & self, at::IntArrayRef kernel_size, at::IntArrayRef stride, at::IntArrayRef padding, at::IntArrayRef dilation, bool ceil_mode) {
  7056. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7057. auto maybe_layer = maybeCurrentDynamicLayer();
  7058. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7059. int64_t cur_level = maybe_layer->layerId();
  7060. if (!isBatchedAtLevel(self, cur_level)) {
  7061. return at::_ops::max_pool1d_with_indices::call(self, kernel_size, stride, padding, dilation, ceil_mode);
  7062. }
  7063. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7064. auto results = batch_rule(self_value, self_bdim, kernel_size, stride, padding, dilation, ceil_mode);
  7065. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  7066. }
  7067. template <typename batch_rule_t, batch_rule_t batch_rule>
  7068. at::Tensor max_pool1d_generated_plumbing(const at::Tensor & self, at::IntArrayRef kernel_size, at::IntArrayRef stride, at::IntArrayRef padding, at::IntArrayRef dilation, bool ceil_mode) {
  7069. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7070. auto maybe_layer = maybeCurrentDynamicLayer();
  7071. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7072. int64_t cur_level = maybe_layer->layerId();
  7073. if (!isBatchedAtLevel(self, cur_level)) {
  7074. return at::_ops::max_pool1d::call(self, kernel_size, stride, padding, dilation, ceil_mode);
  7075. }
  7076. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7077. auto results = batch_rule(self_value, self_bdim, kernel_size, stride, padding, dilation, ceil_mode);
  7078. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7079. }
  7080. template <typename batch_rule_t, batch_rule_t batch_rule>
  7081. at::Tensor max_pool2d_generated_plumbing(const at::Tensor & self, at::IntArrayRef kernel_size, at::IntArrayRef stride, at::IntArrayRef padding, at::IntArrayRef dilation, bool ceil_mode) {
  7082. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7083. auto maybe_layer = maybeCurrentDynamicLayer();
  7084. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7085. int64_t cur_level = maybe_layer->layerId();
  7086. if (!isBatchedAtLevel(self, cur_level)) {
  7087. return at::_ops::max_pool2d::call(self, kernel_size, stride, padding, dilation, ceil_mode);
  7088. }
  7089. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7090. auto results = batch_rule(self_value, self_bdim, kernel_size, stride, padding, dilation, ceil_mode);
  7091. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7092. }
  7093. template <typename batch_rule_t, batch_rule_t batch_rule>
  7094. at::Tensor max_pool2d_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self, at::IntArrayRef kernel_size, at::IntArrayRef stride, at::IntArrayRef padding, at::IntArrayRef dilation, bool ceil_mode) {
  7095. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7096. auto maybe_layer = maybeCurrentDynamicLayer();
  7097. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7098. int64_t cur_level = maybe_layer->layerId();
  7099. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level)) {
  7100. return at::_ops::max_pool2d_backward::call(grad_output, self, kernel_size, stride, padding, dilation, ceil_mode);
  7101. }
  7102. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  7103. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7104. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim, kernel_size, stride, padding, dilation, ceil_mode);
  7105. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7106. }
  7107. template <typename batch_rule_t, batch_rule_t batch_rule>
  7108. at::Tensor mkldnn_max_pool2d_generated_plumbing(const at::Tensor & self, at::IntArrayRef kernel_size, at::IntArrayRef stride, at::IntArrayRef padding, at::IntArrayRef dilation, bool ceil_mode) {
  7109. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7110. auto maybe_layer = maybeCurrentDynamicLayer();
  7111. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7112. int64_t cur_level = maybe_layer->layerId();
  7113. if (!isBatchedAtLevel(self, cur_level)) {
  7114. return at::_ops::mkldnn_max_pool2d::call(self, kernel_size, stride, padding, dilation, ceil_mode);
  7115. }
  7116. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7117. auto results = batch_rule(self_value, self_bdim, kernel_size, stride, padding, dilation, ceil_mode);
  7118. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7119. }
  7120. template <typename batch_rule_t, batch_rule_t batch_rule>
  7121. at::Tensor mkldnn_max_pool2d_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & output, const at::Tensor & input, at::IntArrayRef kernel_size, at::IntArrayRef stride, at::IntArrayRef padding, at::IntArrayRef dilation, bool ceil_mode) {
  7122. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7123. auto maybe_layer = maybeCurrentDynamicLayer();
  7124. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7125. int64_t cur_level = maybe_layer->layerId();
  7126. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(output, cur_level) && !isBatchedAtLevel(input, cur_level)) {
  7127. return at::_ops::mkldnn_max_pool2d_backward::call(grad_output, output, input, kernel_size, stride, padding, dilation, ceil_mode);
  7128. }
  7129. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  7130. auto [output_value, output_bdim] = unwrapTensorAtLevel(output, cur_level);
  7131. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  7132. auto results = batch_rule(grad_output_value, grad_output_bdim, output_value, output_bdim, input_value, input_bdim, kernel_size, stride, padding, dilation, ceil_mode);
  7133. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7134. }
  7135. template <typename batch_rule_t, batch_rule_t batch_rule>
  7136. at::Tensor mkldnn_max_pool3d_generated_plumbing(const at::Tensor & self, at::IntArrayRef kernel_size, at::IntArrayRef stride, at::IntArrayRef padding, at::IntArrayRef dilation, bool ceil_mode) {
  7137. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7138. auto maybe_layer = maybeCurrentDynamicLayer();
  7139. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7140. int64_t cur_level = maybe_layer->layerId();
  7141. if (!isBatchedAtLevel(self, cur_level)) {
  7142. return at::_ops::mkldnn_max_pool3d::call(self, kernel_size, stride, padding, dilation, ceil_mode);
  7143. }
  7144. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7145. auto results = batch_rule(self_value, self_bdim, kernel_size, stride, padding, dilation, ceil_mode);
  7146. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7147. }
  7148. template <typename batch_rule_t, batch_rule_t batch_rule>
  7149. at::Tensor mkldnn_max_pool3d_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & output, const at::Tensor & input, at::IntArrayRef kernel_size, at::IntArrayRef stride, at::IntArrayRef padding, at::IntArrayRef dilation, bool ceil_mode) {
  7150. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7151. auto maybe_layer = maybeCurrentDynamicLayer();
  7152. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7153. int64_t cur_level = maybe_layer->layerId();
  7154. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(output, cur_level) && !isBatchedAtLevel(input, cur_level)) {
  7155. return at::_ops::mkldnn_max_pool3d_backward::call(grad_output, output, input, kernel_size, stride, padding, dilation, ceil_mode);
  7156. }
  7157. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  7158. auto [output_value, output_bdim] = unwrapTensorAtLevel(output, cur_level);
  7159. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  7160. auto results = batch_rule(grad_output_value, grad_output_bdim, output_value, output_bdim, input_value, input_bdim, kernel_size, stride, padding, dilation, ceil_mode);
  7161. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7162. }
  7163. template <typename batch_rule_t, batch_rule_t batch_rule>
  7164. at::Tensor quantized_max_pool1d_generated_plumbing(const at::Tensor & self, at::IntArrayRef kernel_size, at::IntArrayRef stride, at::IntArrayRef padding, at::IntArrayRef dilation, bool ceil_mode) {
  7165. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7166. auto maybe_layer = maybeCurrentDynamicLayer();
  7167. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7168. int64_t cur_level = maybe_layer->layerId();
  7169. if (!isBatchedAtLevel(self, cur_level)) {
  7170. return at::_ops::quantized_max_pool1d::call(self, kernel_size, stride, padding, dilation, ceil_mode);
  7171. }
  7172. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7173. auto results = batch_rule(self_value, self_bdim, kernel_size, stride, padding, dilation, ceil_mode);
  7174. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7175. }
  7176. template <typename batch_rule_t, batch_rule_t batch_rule>
  7177. at::Tensor quantized_max_pool2d_generated_plumbing(const at::Tensor & self, at::IntArrayRef kernel_size, at::IntArrayRef stride, at::IntArrayRef padding, at::IntArrayRef dilation, bool ceil_mode) {
  7178. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7179. auto maybe_layer = maybeCurrentDynamicLayer();
  7180. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7181. int64_t cur_level = maybe_layer->layerId();
  7182. if (!isBatchedAtLevel(self, cur_level)) {
  7183. return at::_ops::quantized_max_pool2d::call(self, kernel_size, stride, padding, dilation, ceil_mode);
  7184. }
  7185. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7186. auto results = batch_rule(self_value, self_bdim, kernel_size, stride, padding, dilation, ceil_mode);
  7187. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7188. }
  7189. template <typename batch_rule_t, batch_rule_t batch_rule>
  7190. at::Tensor quantized_max_pool3d_generated_plumbing(const at::Tensor & self, at::IntArrayRef kernel_size, at::IntArrayRef stride, at::IntArrayRef padding, at::IntArrayRef dilation, bool ceil_mode) {
  7191. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7192. auto maybe_layer = maybeCurrentDynamicLayer();
  7193. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7194. int64_t cur_level = maybe_layer->layerId();
  7195. if (!isBatchedAtLevel(self, cur_level)) {
  7196. return at::_ops::quantized_max_pool3d::call(self, kernel_size, stride, padding, dilation, ceil_mode);
  7197. }
  7198. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7199. auto results = batch_rule(self_value, self_bdim, kernel_size, stride, padding, dilation, ceil_mode);
  7200. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7201. }
  7202. template <typename batch_rule_t, batch_rule_t batch_rule>
  7203. at::Tensor max_pool3d_generated_plumbing(const at::Tensor & self, at::IntArrayRef kernel_size, at::IntArrayRef stride, at::IntArrayRef padding, at::IntArrayRef dilation, bool ceil_mode) {
  7204. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7205. auto maybe_layer = maybeCurrentDynamicLayer();
  7206. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7207. int64_t cur_level = maybe_layer->layerId();
  7208. if (!isBatchedAtLevel(self, cur_level)) {
  7209. return at::_ops::max_pool3d::call(self, kernel_size, stride, padding, dilation, ceil_mode);
  7210. }
  7211. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7212. auto results = batch_rule(self_value, self_bdim, kernel_size, stride, padding, dilation, ceil_mode);
  7213. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7214. }
  7215. template <typename batch_rule_t, batch_rule_t batch_rule>
  7216. at::Tensor mean_generated_plumbing(const at::Tensor & self, ::std::optional<at::ScalarType> dtype) {
  7217. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7218. auto maybe_layer = maybeCurrentDynamicLayer();
  7219. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7220. int64_t cur_level = maybe_layer->layerId();
  7221. if (!isBatchedAtLevel(self, cur_level)) {
  7222. return at::_ops::mean::call(self, dtype);
  7223. }
  7224. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7225. auto results = batch_rule(self_value, self_bdim, dtype);
  7226. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7227. }
  7228. template <typename batch_rule_t, batch_rule_t batch_rule>
  7229. at::Tensor mean_dim_generated_plumbing(const at::Tensor & self, at::OptionalIntArrayRef dim, bool keepdim, ::std::optional<at::ScalarType> dtype) {
  7230. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7231. auto maybe_layer = maybeCurrentDynamicLayer();
  7232. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7233. int64_t cur_level = maybe_layer->layerId();
  7234. if (!isBatchedAtLevel(self, cur_level)) {
  7235. return at::_ops::mean_dim::call(self, dim, keepdim, dtype);
  7236. }
  7237. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7238. auto results = batch_rule(self_value, self_bdim, dim, keepdim, dtype);
  7239. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7240. }
  7241. template <typename batch_rule_t, batch_rule_t batch_rule>
  7242. at::Tensor mean_names_dim_generated_plumbing(const at::Tensor & self, at::DimnameList dim, bool keepdim, ::std::optional<at::ScalarType> dtype) {
  7243. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7244. auto maybe_layer = maybeCurrentDynamicLayer();
  7245. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7246. int64_t cur_level = maybe_layer->layerId();
  7247. if (!isBatchedAtLevel(self, cur_level)) {
  7248. return at::_ops::mean_names_dim::call(self, dim, keepdim, dtype);
  7249. }
  7250. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7251. auto results = batch_rule(self_value, self_bdim, dim, keepdim, dtype);
  7252. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7253. }
  7254. template <typename batch_rule_t, batch_rule_t batch_rule>
  7255. at::Tensor nanmean_generated_plumbing(const at::Tensor & self, at::OptionalIntArrayRef dim, bool keepdim, ::std::optional<at::ScalarType> dtype) {
  7256. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7257. auto maybe_layer = maybeCurrentDynamicLayer();
  7258. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7259. int64_t cur_level = maybe_layer->layerId();
  7260. if (!isBatchedAtLevel(self, cur_level)) {
  7261. return at::_ops::nanmean::call(self, dim, keepdim, dtype);
  7262. }
  7263. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7264. auto results = batch_rule(self_value, self_bdim, dim, keepdim, dtype);
  7265. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7266. }
  7267. template <typename batch_rule_t, batch_rule_t batch_rule>
  7268. at::Tensor median_generated_plumbing(const at::Tensor & self) {
  7269. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7270. auto maybe_layer = maybeCurrentDynamicLayer();
  7271. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7272. int64_t cur_level = maybe_layer->layerId();
  7273. if (!isBatchedAtLevel(self, cur_level)) {
  7274. return at::_ops::median::call(self);
  7275. }
  7276. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7277. auto results = batch_rule(self_value, self_bdim);
  7278. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7279. }
  7280. template <typename batch_rule_t, batch_rule_t batch_rule>
  7281. ::std::tuple<at::Tensor,at::Tensor> median_dim_generated_plumbing(const at::Tensor & self, int64_t dim, bool keepdim) {
  7282. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7283. auto maybe_layer = maybeCurrentDynamicLayer();
  7284. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7285. int64_t cur_level = maybe_layer->layerId();
  7286. if (!isBatchedAtLevel(self, cur_level)) {
  7287. return at::_ops::median_dim::call(self, dim, keepdim);
  7288. }
  7289. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7290. auto results = batch_rule(self_value, self_bdim, dim, keepdim);
  7291. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  7292. }
  7293. template <typename batch_rule_t, batch_rule_t batch_rule>
  7294. ::std::tuple<at::Tensor,at::Tensor> median_names_dim_generated_plumbing(const at::Tensor & self, at::Dimname dim, bool keepdim) {
  7295. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7296. auto maybe_layer = maybeCurrentDynamicLayer();
  7297. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7298. int64_t cur_level = maybe_layer->layerId();
  7299. if (!isBatchedAtLevel(self, cur_level)) {
  7300. return at::_ops::median_names_dim::call(self, dim, keepdim);
  7301. }
  7302. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7303. auto results = batch_rule(self_value, self_bdim, dim, keepdim);
  7304. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  7305. }
  7306. template <typename batch_rule_t, batch_rule_t batch_rule>
  7307. at::Tensor nanmedian_generated_plumbing(const at::Tensor & self) {
  7308. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7309. auto maybe_layer = maybeCurrentDynamicLayer();
  7310. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7311. int64_t cur_level = maybe_layer->layerId();
  7312. if (!isBatchedAtLevel(self, cur_level)) {
  7313. return at::_ops::nanmedian::call(self);
  7314. }
  7315. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7316. auto results = batch_rule(self_value, self_bdim);
  7317. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7318. }
  7319. template <typename batch_rule_t, batch_rule_t batch_rule>
  7320. ::std::tuple<at::Tensor,at::Tensor> nanmedian_dim_generated_plumbing(const at::Tensor & self, int64_t dim, bool keepdim) {
  7321. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7322. auto maybe_layer = maybeCurrentDynamicLayer();
  7323. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7324. int64_t cur_level = maybe_layer->layerId();
  7325. if (!isBatchedAtLevel(self, cur_level)) {
  7326. return at::_ops::nanmedian_dim::call(self, dim, keepdim);
  7327. }
  7328. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7329. auto results = batch_rule(self_value, self_bdim, dim, keepdim);
  7330. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  7331. }
  7332. template <typename batch_rule_t, batch_rule_t batch_rule>
  7333. ::std::tuple<at::Tensor,at::Tensor> nanmedian_names_dim_generated_plumbing(const at::Tensor & self, at::Dimname dim, bool keepdim) {
  7334. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7335. auto maybe_layer = maybeCurrentDynamicLayer();
  7336. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7337. int64_t cur_level = maybe_layer->layerId();
  7338. if (!isBatchedAtLevel(self, cur_level)) {
  7339. return at::_ops::nanmedian_names_dim::call(self, dim, keepdim);
  7340. }
  7341. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7342. auto results = batch_rule(self_value, self_bdim, dim, keepdim);
  7343. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  7344. }
  7345. template <typename batch_rule_t, batch_rule_t batch_rule>
  7346. ::std::tuple<at::Tensor,at::Tensor> min_dim_generated_plumbing(const at::Tensor & self, int64_t dim, bool keepdim) {
  7347. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7348. auto maybe_layer = maybeCurrentDynamicLayer();
  7349. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7350. int64_t cur_level = maybe_layer->layerId();
  7351. if (!isBatchedAtLevel(self, cur_level)) {
  7352. return at::_ops::min_dim::call(self, dim, keepdim);
  7353. }
  7354. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7355. auto results = batch_rule(self_value, self_bdim, dim, keepdim);
  7356. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  7357. }
  7358. template <typename batch_rule_t, batch_rule_t batch_rule>
  7359. ::std::tuple<at::Tensor,at::Tensor> min_names_dim_generated_plumbing(const at::Tensor & self, at::Dimname dim, bool keepdim) {
  7360. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7361. auto maybe_layer = maybeCurrentDynamicLayer();
  7362. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7363. int64_t cur_level = maybe_layer->layerId();
  7364. if (!isBatchedAtLevel(self, cur_level)) {
  7365. return at::_ops::min_names_dim::call(self, dim, keepdim);
  7366. }
  7367. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7368. auto results = batch_rule(self_value, self_bdim, dim, keepdim);
  7369. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  7370. }
  7371. template <typename batch_rule_t, batch_rule_t batch_rule>
  7372. at::Tensor amin_generated_plumbing(const at::Tensor & self, at::IntArrayRef dim, bool keepdim) {
  7373. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7374. auto maybe_layer = maybeCurrentDynamicLayer();
  7375. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7376. int64_t cur_level = maybe_layer->layerId();
  7377. if (!isBatchedAtLevel(self, cur_level)) {
  7378. return at::_ops::amin::call(self, dim, keepdim);
  7379. }
  7380. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7381. auto results = batch_rule(self_value, self_bdim, dim, keepdim);
  7382. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7383. }
  7384. template <typename batch_rule_t, batch_rule_t batch_rule>
  7385. at::Tensor _mps_convolution_generated_plumbing(const at::Tensor & self, const at::Tensor & weight, const ::std::optional<at::Tensor> & bias, c10::SymIntArrayRef padding, c10::SymIntArrayRef stride, c10::SymIntArrayRef dilation, c10::SymInt groups) {
  7386. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7387. auto maybe_layer = maybeCurrentDynamicLayer();
  7388. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7389. int64_t cur_level = maybe_layer->layerId();
  7390. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  7391. return at::_ops::_mps_convolution::call(self, weight, bias, padding, stride, dilation, groups);
  7392. }
  7393. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7394. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  7395. std::optional<Tensor> bias_value;
  7396. std::optional<int64_t> bias_bdim;
  7397. if (bias) {
  7398. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  7399. }
  7400. auto results = batch_rule(self_value, self_bdim, weight_value, weight_bdim, bias_value, bias_bdim, padding, stride, dilation, groups);
  7401. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7402. }
  7403. template <typename batch_rule_t, batch_rule_t batch_rule>
  7404. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> mps_convolution_backward_generated_plumbing(const at::Tensor & self, const at::Tensor & grad_output, const at::Tensor & weight, c10::SymIntArrayRef padding, c10::SymIntArrayRef stride, c10::SymIntArrayRef dilation, c10::SymInt groups, ::std::array<bool,3> output_mask) {
  7405. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7406. auto maybe_layer = maybeCurrentDynamicLayer();
  7407. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7408. int64_t cur_level = maybe_layer->layerId();
  7409. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(weight, cur_level)) {
  7410. return at::_ops::mps_convolution_backward::call(self, grad_output, weight, padding, stride, dilation, groups, output_mask);
  7411. }
  7412. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7413. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  7414. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  7415. auto results = batch_rule(self_value, self_bdim, grad_output_value, grad_output_bdim, weight_value, weight_bdim, padding, stride, dilation, groups, output_mask);
  7416. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  7417. }
  7418. template <typename batch_rule_t, batch_rule_t batch_rule>
  7419. at::Tensor mkldnn_convolution_generated_plumbing(const at::Tensor & self, const at::Tensor & weight, const ::std::optional<at::Tensor> & bias, c10::SymIntArrayRef padding, c10::SymIntArrayRef stride, c10::SymIntArrayRef dilation, c10::SymInt groups) {
  7420. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7421. auto maybe_layer = maybeCurrentDynamicLayer();
  7422. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7423. int64_t cur_level = maybe_layer->layerId();
  7424. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  7425. return at::_ops::mkldnn_convolution::call(self, weight, bias, padding, stride, dilation, groups);
  7426. }
  7427. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7428. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  7429. std::optional<Tensor> bias_value;
  7430. std::optional<int64_t> bias_bdim;
  7431. if (bias) {
  7432. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  7433. }
  7434. auto results = batch_rule(self_value, self_bdim, weight_value, weight_bdim, bias_value, bias_bdim, padding, stride, dilation, groups);
  7435. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7436. }
  7437. template <typename batch_rule_t, batch_rule_t batch_rule>
  7438. ::std::tuple<at::Tensor,at::Tensor,at::Tensor,at::Tensor> mkldnn_rnn_layer_generated_plumbing(const at::Tensor & input, const at::Tensor & weight0, const at::Tensor & weight1, const at::Tensor & weight2, const at::Tensor & weight3, const at::Tensor & hx_, const at::Tensor & cx_, bool reverse, at::IntArrayRef batch_sizes, int64_t mode, int64_t hidden_size, int64_t num_layers, bool has_biases, bool bidirectional, bool batch_first, bool train) {
  7439. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7440. auto maybe_layer = maybeCurrentDynamicLayer();
  7441. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7442. int64_t cur_level = maybe_layer->layerId();
  7443. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight0, cur_level) && !isBatchedAtLevel(weight1, cur_level) && !isBatchedAtLevel(weight2, cur_level) && !isBatchedAtLevel(weight3, cur_level) && !isBatchedAtLevel(hx_, cur_level) && !isBatchedAtLevel(cx_, cur_level)) {
  7444. return at::_ops::mkldnn_rnn_layer::call(input, weight0, weight1, weight2, weight3, hx_, cx_, reverse, batch_sizes, mode, hidden_size, num_layers, has_biases, bidirectional, batch_first, train);
  7445. }
  7446. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  7447. auto [weight0_value, weight0_bdim] = unwrapTensorAtLevel(weight0, cur_level);
  7448. auto [weight1_value, weight1_bdim] = unwrapTensorAtLevel(weight1, cur_level);
  7449. auto [weight2_value, weight2_bdim] = unwrapTensorAtLevel(weight2, cur_level);
  7450. auto [weight3_value, weight3_bdim] = unwrapTensorAtLevel(weight3, cur_level);
  7451. auto [hx__value, hx__bdim] = unwrapTensorAtLevel(hx_, cur_level);
  7452. auto [cx__value, cx__bdim] = unwrapTensorAtLevel(cx_, cur_level);
  7453. auto results = batch_rule(input_value, input_bdim, weight0_value, weight0_bdim, weight1_value, weight1_bdim, weight2_value, weight2_bdim, weight3_value, weight3_bdim, hx__value, hx__bdim, cx__value, cx__bdim, reverse, batch_sizes, mode, hidden_size, num_layers, has_biases, bidirectional, batch_first, train);
  7454. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level), makeBatched(std::get<6>(results), std::get<7>(results), cur_level));
  7455. }
  7456. template <typename batch_rule_t, batch_rule_t batch_rule>
  7457. ::std::tuple<at::Tensor,at::Tensor,at::Tensor,at::Tensor,at::Tensor,at::Tensor,at::Tensor> mkldnn_rnn_layer_backward_generated_plumbing(const at::Tensor & input, const at::Tensor & weight1, const at::Tensor & weight2, const at::Tensor & weight3, const at::Tensor & weight4, const at::Tensor & hx_, const at::Tensor & cx_tmp, const at::Tensor & output, const at::Tensor & hy_, const at::Tensor & cy_, const ::std::optional<at::Tensor> & grad_output, const ::std::optional<at::Tensor> & grad_hy, const ::std::optional<at::Tensor> & grad_cy, bool reverse, int64_t mode, int64_t hidden_size, int64_t num_layers, bool has_biases, bool train, bool bidirectional, at::IntArrayRef batch_sizes, bool batch_first, const at::Tensor & workspace) {
  7458. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7459. auto maybe_layer = maybeCurrentDynamicLayer();
  7460. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7461. int64_t cur_level = maybe_layer->layerId();
  7462. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight1, cur_level) && !isBatchedAtLevel(weight2, cur_level) && !isBatchedAtLevel(weight3, cur_level) && !isBatchedAtLevel(weight4, cur_level) && !isBatchedAtLevel(hx_, cur_level) && !isBatchedAtLevel(cx_tmp, cur_level) && !isBatchedAtLevel(output, cur_level) && !isBatchedAtLevel(hy_, cur_level) && !isBatchedAtLevel(cy_, cur_level) && !isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(grad_hy, cur_level) && !isBatchedAtLevel(grad_cy, cur_level) && !isBatchedAtLevel(workspace, cur_level)) {
  7463. return at::_ops::mkldnn_rnn_layer_backward::call(input, weight1, weight2, weight3, weight4, hx_, cx_tmp, output, hy_, cy_, grad_output, grad_hy, grad_cy, reverse, mode, hidden_size, num_layers, has_biases, train, bidirectional, batch_sizes, batch_first, workspace);
  7464. }
  7465. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  7466. auto [weight1_value, weight1_bdim] = unwrapTensorAtLevel(weight1, cur_level);
  7467. auto [weight2_value, weight2_bdim] = unwrapTensorAtLevel(weight2, cur_level);
  7468. auto [weight3_value, weight3_bdim] = unwrapTensorAtLevel(weight3, cur_level);
  7469. auto [weight4_value, weight4_bdim] = unwrapTensorAtLevel(weight4, cur_level);
  7470. auto [hx__value, hx__bdim] = unwrapTensorAtLevel(hx_, cur_level);
  7471. auto [cx_tmp_value, cx_tmp_bdim] = unwrapTensorAtLevel(cx_tmp, cur_level);
  7472. auto [output_value, output_bdim] = unwrapTensorAtLevel(output, cur_level);
  7473. auto [hy__value, hy__bdim] = unwrapTensorAtLevel(hy_, cur_level);
  7474. auto [cy__value, cy__bdim] = unwrapTensorAtLevel(cy_, cur_level);
  7475. auto [workspace_value, workspace_bdim] = unwrapTensorAtLevel(workspace, cur_level);
  7476. std::optional<Tensor> grad_output_value;
  7477. std::optional<int64_t> grad_output_bdim;
  7478. if (grad_output) {
  7479. std::tie(grad_output_value, grad_output_bdim) = unwrapTensorAtLevel(grad_output.value(), cur_level);
  7480. }
  7481. std::optional<Tensor> grad_hy_value;
  7482. std::optional<int64_t> grad_hy_bdim;
  7483. if (grad_hy) {
  7484. std::tie(grad_hy_value, grad_hy_bdim) = unwrapTensorAtLevel(grad_hy.value(), cur_level);
  7485. }
  7486. std::optional<Tensor> grad_cy_value;
  7487. std::optional<int64_t> grad_cy_bdim;
  7488. if (grad_cy) {
  7489. std::tie(grad_cy_value, grad_cy_bdim) = unwrapTensorAtLevel(grad_cy.value(), cur_level);
  7490. }
  7491. auto results = batch_rule(input_value, input_bdim, weight1_value, weight1_bdim, weight2_value, weight2_bdim, weight3_value, weight3_bdim, weight4_value, weight4_bdim, hx__value, hx__bdim, cx_tmp_value, cx_tmp_bdim, output_value, output_bdim, hy__value, hy__bdim, cy__value, cy__bdim, grad_output_value, grad_output_bdim, grad_hy_value, grad_hy_bdim, grad_cy_value, grad_cy_bdim, reverse, mode, hidden_size, num_layers, has_biases, train, bidirectional, batch_sizes, batch_first, workspace_value, workspace_bdim);
  7492. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level), makeBatched(std::get<6>(results), std::get<7>(results), cur_level), makeBatched(std::get<8>(results), std::get<9>(results), cur_level), makeBatched(std::get<10>(results), std::get<11>(results), cur_level), makeBatched(std::get<12>(results), std::get<13>(results), cur_level));
  7493. }
  7494. template <typename batch_rule_t, batch_rule_t batch_rule>
  7495. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> miopen_batch_norm_generated_plumbing(const at::Tensor & input, const at::Tensor & weight, const ::std::optional<at::Tensor> & bias, const ::std::optional<at::Tensor> & running_mean, const ::std::optional<at::Tensor> & running_var, bool training, double exponential_average_factor, double epsilon) {
  7496. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7497. auto maybe_layer = maybeCurrentDynamicLayer();
  7498. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7499. int64_t cur_level = maybe_layer->layerId();
  7500. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level) && !isBatchedAtLevel(running_mean, cur_level) && !isBatchedAtLevel(running_var, cur_level)) {
  7501. return at::_ops::miopen_batch_norm::call(input, weight, bias, running_mean, running_var, training, exponential_average_factor, epsilon);
  7502. }
  7503. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  7504. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  7505. std::optional<Tensor> bias_value;
  7506. std::optional<int64_t> bias_bdim;
  7507. if (bias) {
  7508. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  7509. }
  7510. std::optional<Tensor> running_mean_value;
  7511. std::optional<int64_t> running_mean_bdim;
  7512. if (running_mean) {
  7513. std::tie(running_mean_value, running_mean_bdim) = unwrapTensorAtLevel(running_mean.value(), cur_level);
  7514. }
  7515. std::optional<Tensor> running_var_value;
  7516. std::optional<int64_t> running_var_bdim;
  7517. if (running_var) {
  7518. std::tie(running_var_value, running_var_bdim) = unwrapTensorAtLevel(running_var.value(), cur_level);
  7519. }
  7520. auto results = batch_rule(input_value, input_bdim, weight_value, weight_bdim, bias_value, bias_bdim, running_mean_value, running_mean_bdim, running_var_value, running_var_bdim, training, exponential_average_factor, epsilon);
  7521. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  7522. }
  7523. template <typename batch_rule_t, batch_rule_t batch_rule>
  7524. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> miopen_batch_norm_backward_generated_plumbing(const at::Tensor & input, const at::Tensor & grad_output, const at::Tensor & weight, const ::std::optional<at::Tensor> & running_mean, const ::std::optional<at::Tensor> & running_var, const ::std::optional<at::Tensor> & save_mean, const ::std::optional<at::Tensor> & save_var, double epsilon) {
  7525. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7526. auto maybe_layer = maybeCurrentDynamicLayer();
  7527. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7528. int64_t cur_level = maybe_layer->layerId();
  7529. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(running_mean, cur_level) && !isBatchedAtLevel(running_var, cur_level) && !isBatchedAtLevel(save_mean, cur_level) && !isBatchedAtLevel(save_var, cur_level)) {
  7530. return at::_ops::miopen_batch_norm_backward::call(input, grad_output, weight, running_mean, running_var, save_mean, save_var, epsilon);
  7531. }
  7532. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  7533. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  7534. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  7535. std::optional<Tensor> running_mean_value;
  7536. std::optional<int64_t> running_mean_bdim;
  7537. if (running_mean) {
  7538. std::tie(running_mean_value, running_mean_bdim) = unwrapTensorAtLevel(running_mean.value(), cur_level);
  7539. }
  7540. std::optional<Tensor> running_var_value;
  7541. std::optional<int64_t> running_var_bdim;
  7542. if (running_var) {
  7543. std::tie(running_var_value, running_var_bdim) = unwrapTensorAtLevel(running_var.value(), cur_level);
  7544. }
  7545. std::optional<Tensor> save_mean_value;
  7546. std::optional<int64_t> save_mean_bdim;
  7547. if (save_mean) {
  7548. std::tie(save_mean_value, save_mean_bdim) = unwrapTensorAtLevel(save_mean.value(), cur_level);
  7549. }
  7550. std::optional<Tensor> save_var_value;
  7551. std::optional<int64_t> save_var_bdim;
  7552. if (save_var) {
  7553. std::tie(save_var_value, save_var_bdim) = unwrapTensorAtLevel(save_var.value(), cur_level);
  7554. }
  7555. auto results = batch_rule(input_value, input_bdim, grad_output_value, grad_output_bdim, weight_value, weight_bdim, running_mean_value, running_mean_bdim, running_var_value, running_var_bdim, save_mean_value, save_mean_bdim, save_var_value, save_var_bdim, epsilon);
  7556. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  7557. }
  7558. template <typename batch_rule_t, batch_rule_t batch_rule>
  7559. at::Tensor miopen_convolution_generated_plumbing(const at::Tensor & self, const at::Tensor & weight, const ::std::optional<at::Tensor> & bias, c10::SymIntArrayRef padding, c10::SymIntArrayRef stride, c10::SymIntArrayRef dilation, c10::SymInt groups, bool benchmark, bool deterministic) {
  7560. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7561. auto maybe_layer = maybeCurrentDynamicLayer();
  7562. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7563. int64_t cur_level = maybe_layer->layerId();
  7564. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  7565. return at::_ops::miopen_convolution::call(self, weight, bias, padding, stride, dilation, groups, benchmark, deterministic);
  7566. }
  7567. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7568. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  7569. std::optional<Tensor> bias_value;
  7570. std::optional<int64_t> bias_bdim;
  7571. if (bias) {
  7572. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  7573. }
  7574. auto results = batch_rule(self_value, self_bdim, weight_value, weight_bdim, bias_value, bias_bdim, padding, stride, dilation, groups, benchmark, deterministic);
  7575. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7576. }
  7577. template <typename batch_rule_t, batch_rule_t batch_rule>
  7578. at::Tensor miopen_convolution_transpose_generated_plumbing(const at::Tensor & self, const at::Tensor & weight, const ::std::optional<at::Tensor> & bias, c10::SymIntArrayRef padding, c10::SymIntArrayRef output_padding, c10::SymIntArrayRef stride, c10::SymIntArrayRef dilation, c10::SymInt groups, bool benchmark, bool deterministic) {
  7579. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7580. auto maybe_layer = maybeCurrentDynamicLayer();
  7581. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7582. int64_t cur_level = maybe_layer->layerId();
  7583. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  7584. return at::_ops::miopen_convolution_transpose::call(self, weight, bias, padding, output_padding, stride, dilation, groups, benchmark, deterministic);
  7585. }
  7586. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7587. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  7588. std::optional<Tensor> bias_value;
  7589. std::optional<int64_t> bias_bdim;
  7590. if (bias) {
  7591. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  7592. }
  7593. auto results = batch_rule(self_value, self_bdim, weight_value, weight_bdim, bias_value, bias_bdim, padding, output_padding, stride, dilation, groups, benchmark, deterministic);
  7594. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7595. }
  7596. template <typename batch_rule_t, batch_rule_t batch_rule>
  7597. at::Tensor miopen_depthwise_convolution_generated_plumbing(const at::Tensor & self, const at::Tensor & weight, const ::std::optional<at::Tensor> & bias, c10::SymIntArrayRef padding, c10::SymIntArrayRef stride, c10::SymIntArrayRef dilation, c10::SymInt groups, bool benchmark, bool deterministic) {
  7598. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7599. auto maybe_layer = maybeCurrentDynamicLayer();
  7600. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7601. int64_t cur_level = maybe_layer->layerId();
  7602. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  7603. return at::_ops::miopen_depthwise_convolution::call(self, weight, bias, padding, stride, dilation, groups, benchmark, deterministic);
  7604. }
  7605. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7606. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  7607. std::optional<Tensor> bias_value;
  7608. std::optional<int64_t> bias_bdim;
  7609. if (bias) {
  7610. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  7611. }
  7612. auto results = batch_rule(self_value, self_bdim, weight_value, weight_bdim, bias_value, bias_bdim, padding, stride, dilation, groups, benchmark, deterministic);
  7613. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7614. }
  7615. template <typename batch_rule_t, batch_rule_t batch_rule>
  7616. at::Tensor miopen_convolution_relu_generated_plumbing(const at::Tensor & self, const at::Tensor & weight, const ::std::optional<at::Tensor> & bias, c10::SymIntArrayRef stride, c10::SymIntArrayRef padding, c10::SymIntArrayRef dilation, c10::SymInt groups) {
  7617. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7618. auto maybe_layer = maybeCurrentDynamicLayer();
  7619. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7620. int64_t cur_level = maybe_layer->layerId();
  7621. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  7622. return at::_ops::miopen_convolution_relu::call(self, weight, bias, stride, padding, dilation, groups);
  7623. }
  7624. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7625. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  7626. std::optional<Tensor> bias_value;
  7627. std::optional<int64_t> bias_bdim;
  7628. if (bias) {
  7629. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  7630. }
  7631. auto results = batch_rule(self_value, self_bdim, weight_value, weight_bdim, bias_value, bias_bdim, stride, padding, dilation, groups);
  7632. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7633. }
  7634. template <typename batch_rule_t, batch_rule_t batch_rule>
  7635. at::Tensor miopen_convolution_add_relu_generated_plumbing(const at::Tensor & self, const at::Tensor & weight, const at::Tensor & z, const ::std::optional<at::Scalar> & alpha, const ::std::optional<at::Tensor> & bias, c10::SymIntArrayRef stride, c10::SymIntArrayRef padding, c10::SymIntArrayRef dilation, c10::SymInt groups) {
  7636. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7637. auto maybe_layer = maybeCurrentDynamicLayer();
  7638. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7639. int64_t cur_level = maybe_layer->layerId();
  7640. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(z, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  7641. return at::_ops::miopen_convolution_add_relu::call(self, weight, z, alpha, bias, stride, padding, dilation, groups);
  7642. }
  7643. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7644. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  7645. auto [z_value, z_bdim] = unwrapTensorAtLevel(z, cur_level);
  7646. std::optional<Tensor> bias_value;
  7647. std::optional<int64_t> bias_bdim;
  7648. if (bias) {
  7649. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  7650. }
  7651. auto results = batch_rule(self_value, self_bdim, weight_value, weight_bdim, z_value, z_bdim, alpha, bias_value, bias_bdim, stride, padding, dilation, groups);
  7652. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7653. }
  7654. template <typename batch_rule_t, batch_rule_t batch_rule>
  7655. ::std::tuple<at::Tensor,at::Tensor,at::Tensor,at::Tensor,at::Tensor> miopen_rnn_generated_plumbing(const at::Tensor & input, at::TensorList weight, int64_t weight_stride0, const at::Tensor & hx, const ::std::optional<at::Tensor> & cx, int64_t mode, int64_t hidden_size, int64_t num_layers, bool batch_first, double dropout, bool train, bool bidirectional, at::IntArrayRef batch_sizes, const ::std::optional<at::Tensor> & dropout_state) {
  7656. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7657. auto maybe_layer = maybeCurrentDynamicLayer();
  7658. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7659. int64_t cur_level = maybe_layer->layerId();
  7660. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(hx, cur_level) && !isBatchedAtLevel(cx, cur_level) && !isBatchedAtLevel(dropout_state, cur_level)) {
  7661. return at::_ops::miopen_rnn::call(input, weight, weight_stride0, hx, cx, mode, hidden_size, num_layers, batch_first, dropout, train, bidirectional, batch_sizes, dropout_state);
  7662. }
  7663. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  7664. auto [hx_value, hx_bdim] = unwrapTensorAtLevel(hx, cur_level);
  7665. std::optional<Tensor> cx_value;
  7666. std::optional<int64_t> cx_bdim;
  7667. if (cx) {
  7668. std::tie(cx_value, cx_bdim) = unwrapTensorAtLevel(cx.value(), cur_level);
  7669. }
  7670. std::optional<Tensor> dropout_state_value;
  7671. std::optional<int64_t> dropout_state_bdim;
  7672. if (dropout_state) {
  7673. std::tie(dropout_state_value, dropout_state_bdim) = unwrapTensorAtLevel(dropout_state.value(), cur_level);
  7674. }
  7675. auto results = batch_rule(input_value, input_bdim, weight, weight_stride0, hx_value, hx_bdim, cx_value, cx_bdim, mode, hidden_size, num_layers, batch_first, dropout, train, bidirectional, batch_sizes, dropout_state_value, dropout_state_bdim);
  7676. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level), makeBatched(std::get<6>(results), std::get<7>(results), cur_level), makeBatched(std::get<8>(results), std::get<9>(results), cur_level));
  7677. }
  7678. template <typename batch_rule_t, batch_rule_t batch_rule>
  7679. ::std::tuple<at::Tensor,at::Tensor,at::Tensor,::std::vector<at::Tensor>> miopen_rnn_backward_generated_plumbing(const at::Tensor & input, at::TensorList weight, int64_t weight_stride0, const at::Tensor & weight_buf, const at::Tensor & hx, const ::std::optional<at::Tensor> & cx, const at::Tensor & output, const ::std::optional<at::Tensor> & grad_output, const ::std::optional<at::Tensor> & grad_hy, const ::std::optional<at::Tensor> & grad_cy, int64_t mode, int64_t hidden_size, int64_t num_layers, bool batch_first, double dropout, bool train, bool bidirectional, at::IntArrayRef batch_sizes, const ::std::optional<at::Tensor> & dropout_state, const at::Tensor & reserve, ::std::array<bool,4> output_mask) {
  7680. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7681. auto maybe_layer = maybeCurrentDynamicLayer();
  7682. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7683. int64_t cur_level = maybe_layer->layerId();
  7684. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(weight_buf, cur_level) && !isBatchedAtLevel(hx, cur_level) && !isBatchedAtLevel(cx, cur_level) && !isBatchedAtLevel(output, cur_level) && !isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(grad_hy, cur_level) && !isBatchedAtLevel(grad_cy, cur_level) && !isBatchedAtLevel(dropout_state, cur_level) && !isBatchedAtLevel(reserve, cur_level)) {
  7685. return at::_ops::miopen_rnn_backward::call(input, weight, weight_stride0, weight_buf, hx, cx, output, grad_output, grad_hy, grad_cy, mode, hidden_size, num_layers, batch_first, dropout, train, bidirectional, batch_sizes, dropout_state, reserve, output_mask);
  7686. }
  7687. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  7688. auto [weight_buf_value, weight_buf_bdim] = unwrapTensorAtLevel(weight_buf, cur_level);
  7689. auto [hx_value, hx_bdim] = unwrapTensorAtLevel(hx, cur_level);
  7690. auto [output_value, output_bdim] = unwrapTensorAtLevel(output, cur_level);
  7691. auto [reserve_value, reserve_bdim] = unwrapTensorAtLevel(reserve, cur_level);
  7692. std::optional<Tensor> cx_value;
  7693. std::optional<int64_t> cx_bdim;
  7694. if (cx) {
  7695. std::tie(cx_value, cx_bdim) = unwrapTensorAtLevel(cx.value(), cur_level);
  7696. }
  7697. std::optional<Tensor> grad_output_value;
  7698. std::optional<int64_t> grad_output_bdim;
  7699. if (grad_output) {
  7700. std::tie(grad_output_value, grad_output_bdim) = unwrapTensorAtLevel(grad_output.value(), cur_level);
  7701. }
  7702. std::optional<Tensor> grad_hy_value;
  7703. std::optional<int64_t> grad_hy_bdim;
  7704. if (grad_hy) {
  7705. std::tie(grad_hy_value, grad_hy_bdim) = unwrapTensorAtLevel(grad_hy.value(), cur_level);
  7706. }
  7707. std::optional<Tensor> grad_cy_value;
  7708. std::optional<int64_t> grad_cy_bdim;
  7709. if (grad_cy) {
  7710. std::tie(grad_cy_value, grad_cy_bdim) = unwrapTensorAtLevel(grad_cy.value(), cur_level);
  7711. }
  7712. std::optional<Tensor> dropout_state_value;
  7713. std::optional<int64_t> dropout_state_bdim;
  7714. if (dropout_state) {
  7715. std::tie(dropout_state_value, dropout_state_bdim) = unwrapTensorAtLevel(dropout_state.value(), cur_level);
  7716. }
  7717. auto results = batch_rule(input_value, input_bdim, weight, weight_stride0, weight_buf_value, weight_buf_bdim, hx_value, hx_bdim, cx_value, cx_bdim, output_value, output_bdim, grad_output_value, grad_output_bdim, grad_hy_value, grad_hy_bdim, grad_cy_value, grad_cy_bdim, mode, hidden_size, num_layers, batch_first, dropout, train, bidirectional, batch_sizes, dropout_state_value, dropout_state_bdim, reserve_value, reserve_bdim, output_mask);
  7718. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level), makeBatchedVector(std::get<6>(results), std::get<7>(results), cur_level));
  7719. }
  7720. template <typename batch_rule_t, batch_rule_t batch_rule>
  7721. at::Tensor mm_generated_plumbing(const at::Tensor & self, const at::Tensor & mat2) {
  7722. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7723. auto maybe_layer = maybeCurrentDynamicLayer();
  7724. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7725. int64_t cur_level = maybe_layer->layerId();
  7726. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(mat2, cur_level)) {
  7727. return at::_ops::mm::call(self, mat2);
  7728. }
  7729. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7730. auto [mat2_value, mat2_bdim] = unwrapTensorAtLevel(mat2, cur_level);
  7731. auto results = batch_rule(self_value, self_bdim, mat2_value, mat2_bdim);
  7732. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7733. }
  7734. template <typename batch_rule_t, batch_rule_t batch_rule>
  7735. at::Tensor mm_dtype_generated_plumbing(const at::Tensor & self, const at::Tensor & mat2, at::ScalarType out_dtype) {
  7736. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7737. auto maybe_layer = maybeCurrentDynamicLayer();
  7738. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7739. int64_t cur_level = maybe_layer->layerId();
  7740. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(mat2, cur_level)) {
  7741. return at::_ops::mm_dtype::call(self, mat2, out_dtype);
  7742. }
  7743. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7744. auto [mat2_value, mat2_bdim] = unwrapTensorAtLevel(mat2, cur_level);
  7745. auto results = batch_rule(self_value, self_bdim, mat2_value, mat2_bdim, out_dtype);
  7746. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7747. }
  7748. template <typename batch_rule_t, batch_rule_t batch_rule>
  7749. at::Tensor _int_mm_generated_plumbing(const at::Tensor & self, const at::Tensor & mat2) {
  7750. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7751. auto maybe_layer = maybeCurrentDynamicLayer();
  7752. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7753. int64_t cur_level = maybe_layer->layerId();
  7754. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(mat2, cur_level)) {
  7755. return at::_ops::_int_mm::call(self, mat2);
  7756. }
  7757. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7758. auto [mat2_value, mat2_bdim] = unwrapTensorAtLevel(mat2, cur_level);
  7759. auto results = batch_rule(self_value, self_bdim, mat2_value, mat2_bdim);
  7760. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7761. }
  7762. template <typename batch_rule_t, batch_rule_t batch_rule>
  7763. at::Tensor _convert_weight_to_int4pack_generated_plumbing(const at::Tensor & self, int64_t innerKTiles) {
  7764. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7765. auto maybe_layer = maybeCurrentDynamicLayer();
  7766. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7767. int64_t cur_level = maybe_layer->layerId();
  7768. if (!isBatchedAtLevel(self, cur_level)) {
  7769. return at::_ops::_convert_weight_to_int4pack::call(self, innerKTiles);
  7770. }
  7771. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7772. auto results = batch_rule(self_value, self_bdim, innerKTiles);
  7773. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7774. }
  7775. template <typename batch_rule_t, batch_rule_t batch_rule>
  7776. at::Tensor _weight_int4pack_mm_generated_plumbing(const at::Tensor & self, const at::Tensor & mat2, int64_t qGroupSize, const at::Tensor & qScaleAndZeros) {
  7777. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7778. auto maybe_layer = maybeCurrentDynamicLayer();
  7779. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7780. int64_t cur_level = maybe_layer->layerId();
  7781. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(mat2, cur_level) && !isBatchedAtLevel(qScaleAndZeros, cur_level)) {
  7782. return at::_ops::_weight_int4pack_mm::call(self, mat2, qGroupSize, qScaleAndZeros);
  7783. }
  7784. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7785. auto [mat2_value, mat2_bdim] = unwrapTensorAtLevel(mat2, cur_level);
  7786. auto [qScaleAndZeros_value, qScaleAndZeros_bdim] = unwrapTensorAtLevel(qScaleAndZeros, cur_level);
  7787. auto results = batch_rule(self_value, self_bdim, mat2_value, mat2_bdim, qGroupSize, qScaleAndZeros_value, qScaleAndZeros_bdim);
  7788. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7789. }
  7790. template <typename batch_rule_t, batch_rule_t batch_rule>
  7791. at::Tensor _weight_int4pack_mm_with_scales_and_zeros_generated_plumbing(const at::Tensor & self, const at::Tensor & mat2, int64_t qGroupSize, const at::Tensor & qScale, const at::Tensor & qZeros) {
  7792. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7793. auto maybe_layer = maybeCurrentDynamicLayer();
  7794. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7795. int64_t cur_level = maybe_layer->layerId();
  7796. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(mat2, cur_level) && !isBatchedAtLevel(qScale, cur_level) && !isBatchedAtLevel(qZeros, cur_level)) {
  7797. return at::_ops::_weight_int4pack_mm_with_scales_and_zeros::call(self, mat2, qGroupSize, qScale, qZeros);
  7798. }
  7799. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7800. auto [mat2_value, mat2_bdim] = unwrapTensorAtLevel(mat2, cur_level);
  7801. auto [qScale_value, qScale_bdim] = unwrapTensorAtLevel(qScale, cur_level);
  7802. auto [qZeros_value, qZeros_bdim] = unwrapTensorAtLevel(qZeros, cur_level);
  7803. auto results = batch_rule(self_value, self_bdim, mat2_value, mat2_bdim, qGroupSize, qScale_value, qScale_bdim, qZeros_value, qZeros_bdim);
  7804. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7805. }
  7806. template <typename batch_rule_t, batch_rule_t batch_rule>
  7807. at::Tensor _convert_weight_to_int4pack_for_cpu_generated_plumbing(const at::Tensor & self, int64_t innerKTiles) {
  7808. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7809. auto maybe_layer = maybeCurrentDynamicLayer();
  7810. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7811. int64_t cur_level = maybe_layer->layerId();
  7812. if (!isBatchedAtLevel(self, cur_level)) {
  7813. return at::_ops::_convert_weight_to_int4pack_for_cpu::call(self, innerKTiles);
  7814. }
  7815. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7816. auto results = batch_rule(self_value, self_bdim, innerKTiles);
  7817. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7818. }
  7819. template <typename batch_rule_t, batch_rule_t batch_rule>
  7820. at::Tensor _weight_int4pack_mm_for_cpu_generated_plumbing(const at::Tensor & self, const at::Tensor & mat2, int64_t qGroupSize, const at::Tensor & qScaleAndZeros) {
  7821. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7822. auto maybe_layer = maybeCurrentDynamicLayer();
  7823. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7824. int64_t cur_level = maybe_layer->layerId();
  7825. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(mat2, cur_level) && !isBatchedAtLevel(qScaleAndZeros, cur_level)) {
  7826. return at::_ops::_weight_int4pack_mm_for_cpu::call(self, mat2, qGroupSize, qScaleAndZeros);
  7827. }
  7828. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7829. auto [mat2_value, mat2_bdim] = unwrapTensorAtLevel(mat2, cur_level);
  7830. auto [qScaleAndZeros_value, qScaleAndZeros_bdim] = unwrapTensorAtLevel(qScaleAndZeros, cur_level);
  7831. auto results = batch_rule(self_value, self_bdim, mat2_value, mat2_bdim, qGroupSize, qScaleAndZeros_value, qScaleAndZeros_bdim);
  7832. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7833. }
  7834. template <typename batch_rule_t, batch_rule_t batch_rule>
  7835. at::Tensor _dyn_quant_pack_4bit_weight_generated_plumbing(const at::Tensor & weights, const at::Tensor & scales_zeros, const ::std::optional<at::Tensor> & bias, int64_t block_size, int64_t in_features, int64_t out_features) {
  7836. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7837. auto maybe_layer = maybeCurrentDynamicLayer();
  7838. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7839. int64_t cur_level = maybe_layer->layerId();
  7840. if (!isBatchedAtLevel(weights, cur_level) && !isBatchedAtLevel(scales_zeros, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  7841. return at::_ops::_dyn_quant_pack_4bit_weight::call(weights, scales_zeros, bias, block_size, in_features, out_features);
  7842. }
  7843. auto [weights_value, weights_bdim] = unwrapTensorAtLevel(weights, cur_level);
  7844. auto [scales_zeros_value, scales_zeros_bdim] = unwrapTensorAtLevel(scales_zeros, cur_level);
  7845. std::optional<Tensor> bias_value;
  7846. std::optional<int64_t> bias_bdim;
  7847. if (bias) {
  7848. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  7849. }
  7850. auto results = batch_rule(weights_value, weights_bdim, scales_zeros_value, scales_zeros_bdim, bias_value, bias_bdim, block_size, in_features, out_features);
  7851. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7852. }
  7853. template <typename batch_rule_t, batch_rule_t batch_rule>
  7854. at::Tensor _dyn_quant_matmul_4bit_generated_plumbing(const at::Tensor & inp, const at::Tensor & packed_weights, int64_t block_size, int64_t in_features, int64_t out_features) {
  7855. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7856. auto maybe_layer = maybeCurrentDynamicLayer();
  7857. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7858. int64_t cur_level = maybe_layer->layerId();
  7859. if (!isBatchedAtLevel(inp, cur_level) && !isBatchedAtLevel(packed_weights, cur_level)) {
  7860. return at::_ops::_dyn_quant_matmul_4bit::call(inp, packed_weights, block_size, in_features, out_features);
  7861. }
  7862. auto [inp_value, inp_bdim] = unwrapTensorAtLevel(inp, cur_level);
  7863. auto [packed_weights_value, packed_weights_bdim] = unwrapTensorAtLevel(packed_weights, cur_level);
  7864. auto results = batch_rule(inp_value, inp_bdim, packed_weights_value, packed_weights_bdim, block_size, in_features, out_features);
  7865. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7866. }
  7867. template <typename batch_rule_t, batch_rule_t batch_rule>
  7868. at::Tensor _weight_int8pack_mm_generated_plumbing(const at::Tensor & self, const at::Tensor & mat2, const at::Tensor & scales) {
  7869. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7870. auto maybe_layer = maybeCurrentDynamicLayer();
  7871. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7872. int64_t cur_level = maybe_layer->layerId();
  7873. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(mat2, cur_level) && !isBatchedAtLevel(scales, cur_level)) {
  7874. return at::_ops::_weight_int8pack_mm::call(self, mat2, scales);
  7875. }
  7876. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7877. auto [mat2_value, mat2_bdim] = unwrapTensorAtLevel(mat2, cur_level);
  7878. auto [scales_value, scales_bdim] = unwrapTensorAtLevel(scales, cur_level);
  7879. auto results = batch_rule(self_value, self_bdim, mat2_value, mat2_bdim, scales_value, scales_bdim);
  7880. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7881. }
  7882. template <typename batch_rule_t, batch_rule_t batch_rule>
  7883. at::Tensor _sparse_mm_generated_plumbing(const at::Tensor & sparse, const at::Tensor & dense) {
  7884. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7885. auto maybe_layer = maybeCurrentDynamicLayer();
  7886. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7887. int64_t cur_level = maybe_layer->layerId();
  7888. if (!isBatchedAtLevel(sparse, cur_level) && !isBatchedAtLevel(dense, cur_level)) {
  7889. return at::_ops::_sparse_mm::call(sparse, dense);
  7890. }
  7891. auto [sparse_value, sparse_bdim] = unwrapTensorAtLevel(sparse, cur_level);
  7892. auto [dense_value, dense_bdim] = unwrapTensorAtLevel(dense, cur_level);
  7893. auto results = batch_rule(sparse_value, sparse_bdim, dense_value, dense_bdim);
  7894. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7895. }
  7896. template <typename batch_rule_t, batch_rule_t batch_rule>
  7897. at::Tensor _sparse_mm_reduce_generated_plumbing(const at::Tensor & sparse, const at::Tensor & dense, c10::string_view reduce) {
  7898. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7899. auto maybe_layer = maybeCurrentDynamicLayer();
  7900. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7901. int64_t cur_level = maybe_layer->layerId();
  7902. if (!isBatchedAtLevel(sparse, cur_level) && !isBatchedAtLevel(dense, cur_level)) {
  7903. return at::_ops::_sparse_mm_reduce::call(sparse, dense, reduce);
  7904. }
  7905. auto [sparse_value, sparse_bdim] = unwrapTensorAtLevel(sparse, cur_level);
  7906. auto [dense_value, dense_bdim] = unwrapTensorAtLevel(dense, cur_level);
  7907. auto results = batch_rule(sparse_value, sparse_bdim, dense_value, dense_bdim, reduce);
  7908. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7909. }
  7910. template <typename batch_rule_t, batch_rule_t batch_rule>
  7911. at::Tensor _sparse_sparse_matmul_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  7912. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7913. auto maybe_layer = maybeCurrentDynamicLayer();
  7914. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7915. int64_t cur_level = maybe_layer->layerId();
  7916. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  7917. return at::_ops::_sparse_sparse_matmul::call(self, other);
  7918. }
  7919. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7920. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  7921. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  7922. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7923. }
  7924. template <typename batch_rule_t, batch_rule_t batch_rule>
  7925. ::std::tuple<at::Tensor,at::Tensor> mode_generated_plumbing(const at::Tensor & self, int64_t dim, bool keepdim) {
  7926. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7927. auto maybe_layer = maybeCurrentDynamicLayer();
  7928. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7929. int64_t cur_level = maybe_layer->layerId();
  7930. if (!isBatchedAtLevel(self, cur_level)) {
  7931. return at::_ops::mode::call(self, dim, keepdim);
  7932. }
  7933. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7934. auto results = batch_rule(self_value, self_bdim, dim, keepdim);
  7935. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  7936. }
  7937. template <typename batch_rule_t, batch_rule_t batch_rule>
  7938. ::std::tuple<at::Tensor,at::Tensor> mode_dimname_generated_plumbing(const at::Tensor & self, at::Dimname dim, bool keepdim) {
  7939. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7940. auto maybe_layer = maybeCurrentDynamicLayer();
  7941. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7942. int64_t cur_level = maybe_layer->layerId();
  7943. if (!isBatchedAtLevel(self, cur_level)) {
  7944. return at::_ops::mode_dimname::call(self, dim, keepdim);
  7945. }
  7946. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7947. auto results = batch_rule(self_value, self_bdim, dim, keepdim);
  7948. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  7949. }
  7950. template <typename batch_rule_t, batch_rule_t batch_rule>
  7951. at::Tensor mul_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  7952. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7953. auto maybe_layer = maybeCurrentDynamicLayer();
  7954. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7955. int64_t cur_level = maybe_layer->layerId();
  7956. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  7957. return at::_ops::mul_Tensor::call(self, other);
  7958. }
  7959. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7960. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  7961. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  7962. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7963. }
  7964. template <typename batch_rule_t, batch_rule_t batch_rule>
  7965. at::Tensor & mul__Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  7966. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7967. auto maybe_layer = maybeCurrentDynamicLayer();
  7968. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  7969. int64_t cur_level = maybe_layer->layerId();
  7970. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  7971. return at::_ops::mul__Tensor::call(self, other);
  7972. }
  7973. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7974. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  7975. batch_rule(self_value, self_bdim, other_value, other_bdim);
  7976. return self;
  7977. }
  7978. template <typename batch_rule_t, batch_rule_t batch_rule>
  7979. at::Tensor mul_Scalar_generated_plumbing(const at::Tensor & self, const at::Scalar & other) {
  7980. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7981. auto maybe_layer = maybeCurrentDynamicLayer();
  7982. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  7983. int64_t cur_level = maybe_layer->layerId();
  7984. if (!isBatchedAtLevel(self, cur_level)) {
  7985. return at::_ops::mul_Scalar::call(self, other);
  7986. }
  7987. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  7988. auto results = batch_rule(self_value, self_bdim, other);
  7989. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  7990. }
  7991. template <typename batch_rule_t, batch_rule_t batch_rule>
  7992. at::Tensor & mul__Scalar_generated_plumbing(at::Tensor & self, const at::Scalar & other) {
  7993. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  7994. auto maybe_layer = maybeCurrentDynamicLayer();
  7995. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  7996. int64_t cur_level = maybe_layer->layerId();
  7997. if (!isBatchedAtLevel(self, cur_level)) {
  7998. return at::_ops::mul__Scalar::call(self, other);
  7999. }
  8000. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8001. batch_rule(self_value, self_bdim, other);
  8002. return self;
  8003. }
  8004. template <typename batch_rule_t, batch_rule_t batch_rule>
  8005. at::Tensor multiply_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  8006. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8007. auto maybe_layer = maybeCurrentDynamicLayer();
  8008. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8009. int64_t cur_level = maybe_layer->layerId();
  8010. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  8011. return at::_ops::multiply_Tensor::call(self, other);
  8012. }
  8013. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8014. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  8015. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  8016. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8017. }
  8018. template <typename batch_rule_t, batch_rule_t batch_rule>
  8019. at::Tensor & multiply__Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  8020. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8021. auto maybe_layer = maybeCurrentDynamicLayer();
  8022. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  8023. int64_t cur_level = maybe_layer->layerId();
  8024. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  8025. return at::_ops::multiply__Tensor::call(self, other);
  8026. }
  8027. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8028. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  8029. batch_rule(self_value, self_bdim, other_value, other_bdim);
  8030. return self;
  8031. }
  8032. template <typename batch_rule_t, batch_rule_t batch_rule>
  8033. at::Tensor multiply_Scalar_generated_plumbing(const at::Tensor & self, const at::Scalar & other) {
  8034. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8035. auto maybe_layer = maybeCurrentDynamicLayer();
  8036. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8037. int64_t cur_level = maybe_layer->layerId();
  8038. if (!isBatchedAtLevel(self, cur_level)) {
  8039. return at::_ops::multiply_Scalar::call(self, other);
  8040. }
  8041. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8042. auto results = batch_rule(self_value, self_bdim, other);
  8043. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8044. }
  8045. template <typename batch_rule_t, batch_rule_t batch_rule>
  8046. at::Tensor & multiply__Scalar_generated_plumbing(at::Tensor & self, const at::Scalar & other) {
  8047. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8048. auto maybe_layer = maybeCurrentDynamicLayer();
  8049. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  8050. int64_t cur_level = maybe_layer->layerId();
  8051. if (!isBatchedAtLevel(self, cur_level)) {
  8052. return at::_ops::multiply__Scalar::call(self, other);
  8053. }
  8054. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8055. batch_rule(self_value, self_bdim, other);
  8056. return self;
  8057. }
  8058. template <typename batch_rule_t, batch_rule_t batch_rule>
  8059. at::Tensor mv_generated_plumbing(const at::Tensor & self, const at::Tensor & vec) {
  8060. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8061. auto maybe_layer = maybeCurrentDynamicLayer();
  8062. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8063. int64_t cur_level = maybe_layer->layerId();
  8064. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(vec, cur_level)) {
  8065. return at::_ops::mv::call(self, vec);
  8066. }
  8067. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8068. auto [vec_value, vec_bdim] = unwrapTensorAtLevel(vec, cur_level);
  8069. auto results = batch_rule(self_value, self_bdim, vec_value, vec_bdim);
  8070. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8071. }
  8072. template <typename batch_rule_t, batch_rule_t batch_rule>
  8073. at::Tensor mvlgamma_generated_plumbing(const at::Tensor & self, int64_t p) {
  8074. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8075. auto maybe_layer = maybeCurrentDynamicLayer();
  8076. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8077. int64_t cur_level = maybe_layer->layerId();
  8078. if (!isBatchedAtLevel(self, cur_level)) {
  8079. return at::_ops::mvlgamma::call(self, p);
  8080. }
  8081. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8082. auto results = batch_rule(self_value, self_bdim, p);
  8083. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8084. }
  8085. template <typename batch_rule_t, batch_rule_t batch_rule>
  8086. at::Tensor & mvlgamma__generated_plumbing(at::Tensor & self, int64_t p) {
  8087. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8088. auto maybe_layer = maybeCurrentDynamicLayer();
  8089. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  8090. int64_t cur_level = maybe_layer->layerId();
  8091. if (!isBatchedAtLevel(self, cur_level)) {
  8092. return at::_ops::mvlgamma_::call(self, p);
  8093. }
  8094. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8095. batch_rule(self_value, self_bdim, p);
  8096. return self;
  8097. }
  8098. template <typename batch_rule_t, batch_rule_t batch_rule>
  8099. at::Tensor narrow_copy_generated_plumbing(const at::Tensor & self, int64_t dim, c10::SymInt start, c10::SymInt length) {
  8100. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8101. auto maybe_layer = maybeCurrentDynamicLayer();
  8102. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8103. int64_t cur_level = maybe_layer->layerId();
  8104. if (!isBatchedAtLevel(self, cur_level)) {
  8105. return at::_ops::narrow_copy::call(self, dim, start, length);
  8106. }
  8107. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8108. auto results = batch_rule(self_value, self_bdim, dim, start, length);
  8109. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8110. }
  8111. template <typename batch_rule_t, batch_rule_t batch_rule>
  8112. at::Tensor narrow_generated_plumbing(const at::Tensor & self, int64_t dim, c10::SymInt start, c10::SymInt length) {
  8113. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8114. auto maybe_layer = maybeCurrentDynamicLayer();
  8115. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8116. int64_t cur_level = maybe_layer->layerId();
  8117. if (!isBatchedAtLevel(self, cur_level)) {
  8118. return at::_ops::narrow::call(self, dim, start, length);
  8119. }
  8120. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8121. auto results = batch_rule(self_value, self_bdim, dim, start, length);
  8122. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8123. }
  8124. template <typename batch_rule_t, batch_rule_t batch_rule>
  8125. at::Tensor narrow_Tensor_generated_plumbing(const at::Tensor & self, int64_t dim, const at::Tensor & start, c10::SymInt length) {
  8126. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8127. auto maybe_layer = maybeCurrentDynamicLayer();
  8128. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8129. int64_t cur_level = maybe_layer->layerId();
  8130. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(start, cur_level)) {
  8131. return at::_ops::narrow_Tensor::call(self, dim, start, length);
  8132. }
  8133. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8134. auto [start_value, start_bdim] = unwrapTensorAtLevel(start, cur_level);
  8135. auto results = batch_rule(self_value, self_bdim, dim, start_value, start_bdim, length);
  8136. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8137. }
  8138. template <typename batch_rule_t, batch_rule_t batch_rule>
  8139. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> native_batch_norm_generated_plumbing(const at::Tensor & input, const ::std::optional<at::Tensor> & weight, const ::std::optional<at::Tensor> & bias, const ::std::optional<at::Tensor> & running_mean, const ::std::optional<at::Tensor> & running_var, bool training, double momentum, double eps) {
  8140. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8141. auto maybe_layer = maybeCurrentDynamicLayer();
  8142. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8143. int64_t cur_level = maybe_layer->layerId();
  8144. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level) && !isBatchedAtLevel(running_mean, cur_level) && !isBatchedAtLevel(running_var, cur_level)) {
  8145. return at::_ops::native_batch_norm::call(input, weight, bias, running_mean, running_var, training, momentum, eps);
  8146. }
  8147. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  8148. std::optional<Tensor> weight_value;
  8149. std::optional<int64_t> weight_bdim;
  8150. if (weight) {
  8151. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  8152. }
  8153. std::optional<Tensor> bias_value;
  8154. std::optional<int64_t> bias_bdim;
  8155. if (bias) {
  8156. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  8157. }
  8158. std::optional<Tensor> running_mean_value;
  8159. std::optional<int64_t> running_mean_bdim;
  8160. if (running_mean) {
  8161. std::tie(running_mean_value, running_mean_bdim) = unwrapTensorAtLevel(running_mean.value(), cur_level);
  8162. }
  8163. std::optional<Tensor> running_var_value;
  8164. std::optional<int64_t> running_var_bdim;
  8165. if (running_var) {
  8166. std::tie(running_var_value, running_var_bdim) = unwrapTensorAtLevel(running_var.value(), cur_level);
  8167. }
  8168. auto results = batch_rule(input_value, input_bdim, weight_value, weight_bdim, bias_value, bias_bdim, running_mean_value, running_mean_bdim, running_var_value, running_var_bdim, training, momentum, eps);
  8169. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  8170. }
  8171. template <typename batch_rule_t, batch_rule_t batch_rule>
  8172. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> _native_batch_norm_legit_no_training_generated_plumbing(const at::Tensor & input, const ::std::optional<at::Tensor> & weight, const ::std::optional<at::Tensor> & bias, const at::Tensor & running_mean, const at::Tensor & running_var, double momentum, double eps) {
  8173. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8174. auto maybe_layer = maybeCurrentDynamicLayer();
  8175. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8176. int64_t cur_level = maybe_layer->layerId();
  8177. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level) && !isBatchedAtLevel(running_mean, cur_level) && !isBatchedAtLevel(running_var, cur_level)) {
  8178. return at::_ops::_native_batch_norm_legit_no_training::call(input, weight, bias, running_mean, running_var, momentum, eps);
  8179. }
  8180. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  8181. auto [running_mean_value, running_mean_bdim] = unwrapTensorAtLevel(running_mean, cur_level);
  8182. auto [running_var_value, running_var_bdim] = unwrapTensorAtLevel(running_var, cur_level);
  8183. std::optional<Tensor> weight_value;
  8184. std::optional<int64_t> weight_bdim;
  8185. if (weight) {
  8186. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  8187. }
  8188. std::optional<Tensor> bias_value;
  8189. std::optional<int64_t> bias_bdim;
  8190. if (bias) {
  8191. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  8192. }
  8193. auto results = batch_rule(input_value, input_bdim, weight_value, weight_bdim, bias_value, bias_bdim, running_mean_value, running_mean_bdim, running_var_value, running_var_bdim, momentum, eps);
  8194. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  8195. }
  8196. template <typename batch_rule_t, batch_rule_t batch_rule>
  8197. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> _native_batch_norm_legit_no_stats_generated_plumbing(const at::Tensor & input, const ::std::optional<at::Tensor> & weight, const ::std::optional<at::Tensor> & bias, bool training, double momentum, double eps) {
  8198. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8199. auto maybe_layer = maybeCurrentDynamicLayer();
  8200. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8201. int64_t cur_level = maybe_layer->layerId();
  8202. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  8203. return at::_ops::_native_batch_norm_legit_no_stats::call(input, weight, bias, training, momentum, eps);
  8204. }
  8205. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  8206. std::optional<Tensor> weight_value;
  8207. std::optional<int64_t> weight_bdim;
  8208. if (weight) {
  8209. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  8210. }
  8211. std::optional<Tensor> bias_value;
  8212. std::optional<int64_t> bias_bdim;
  8213. if (bias) {
  8214. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  8215. }
  8216. auto results = batch_rule(input_value, input_bdim, weight_value, weight_bdim, bias_value, bias_bdim, training, momentum, eps);
  8217. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  8218. }
  8219. template <typename batch_rule_t, batch_rule_t batch_rule>
  8220. ::std::tuple<at::Tensor,at::Tensor> batch_norm_stats_generated_plumbing(const at::Tensor & input, double eps) {
  8221. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8222. auto maybe_layer = maybeCurrentDynamicLayer();
  8223. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8224. int64_t cur_level = maybe_layer->layerId();
  8225. if (!isBatchedAtLevel(input, cur_level)) {
  8226. return at::_ops::batch_norm_stats::call(input, eps);
  8227. }
  8228. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  8229. auto results = batch_rule(input_value, input_bdim, eps);
  8230. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  8231. }
  8232. template <typename batch_rule_t, batch_rule_t batch_rule>
  8233. at::Tensor batch_norm_elemt_generated_plumbing(const at::Tensor & input, const ::std::optional<at::Tensor> & weight, const ::std::optional<at::Tensor> & bias, const at::Tensor & mean, const at::Tensor & invstd, double eps) {
  8234. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8235. auto maybe_layer = maybeCurrentDynamicLayer();
  8236. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8237. int64_t cur_level = maybe_layer->layerId();
  8238. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level) && !isBatchedAtLevel(mean, cur_level) && !isBatchedAtLevel(invstd, cur_level)) {
  8239. return at::_ops::batch_norm_elemt::call(input, weight, bias, mean, invstd, eps);
  8240. }
  8241. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  8242. auto [mean_value, mean_bdim] = unwrapTensorAtLevel(mean, cur_level);
  8243. auto [invstd_value, invstd_bdim] = unwrapTensorAtLevel(invstd, cur_level);
  8244. std::optional<Tensor> weight_value;
  8245. std::optional<int64_t> weight_bdim;
  8246. if (weight) {
  8247. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  8248. }
  8249. std::optional<Tensor> bias_value;
  8250. std::optional<int64_t> bias_bdim;
  8251. if (bias) {
  8252. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  8253. }
  8254. auto results = batch_rule(input_value, input_bdim, weight_value, weight_bdim, bias_value, bias_bdim, mean_value, mean_bdim, invstd_value, invstd_bdim, eps);
  8255. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8256. }
  8257. template <typename batch_rule_t, batch_rule_t batch_rule>
  8258. ::std::tuple<at::Tensor,at::Tensor> batch_norm_gather_stats_generated_plumbing(const at::Tensor & input, const at::Tensor & mean, const at::Tensor & invstd, const ::std::optional<at::Tensor> & running_mean, const ::std::optional<at::Tensor> & running_var, double momentum, double eps, int64_t count) {
  8259. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8260. auto maybe_layer = maybeCurrentDynamicLayer();
  8261. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8262. int64_t cur_level = maybe_layer->layerId();
  8263. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(mean, cur_level) && !isBatchedAtLevel(invstd, cur_level) && !isBatchedAtLevel(running_mean, cur_level) && !isBatchedAtLevel(running_var, cur_level)) {
  8264. return at::_ops::batch_norm_gather_stats::call(input, mean, invstd, running_mean, running_var, momentum, eps, count);
  8265. }
  8266. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  8267. auto [mean_value, mean_bdim] = unwrapTensorAtLevel(mean, cur_level);
  8268. auto [invstd_value, invstd_bdim] = unwrapTensorAtLevel(invstd, cur_level);
  8269. std::optional<Tensor> running_mean_value;
  8270. std::optional<int64_t> running_mean_bdim;
  8271. if (running_mean) {
  8272. std::tie(running_mean_value, running_mean_bdim) = unwrapTensorAtLevel(running_mean.value(), cur_level);
  8273. }
  8274. std::optional<Tensor> running_var_value;
  8275. std::optional<int64_t> running_var_bdim;
  8276. if (running_var) {
  8277. std::tie(running_var_value, running_var_bdim) = unwrapTensorAtLevel(running_var.value(), cur_level);
  8278. }
  8279. auto results = batch_rule(input_value, input_bdim, mean_value, mean_bdim, invstd_value, invstd_bdim, running_mean_value, running_mean_bdim, running_var_value, running_var_bdim, momentum, eps, count);
  8280. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  8281. }
  8282. template <typename batch_rule_t, batch_rule_t batch_rule>
  8283. ::std::tuple<at::Tensor,at::Tensor> batch_norm_gather_stats_with_counts_generated_plumbing(const at::Tensor & input, const at::Tensor & mean, const at::Tensor & invstd, const ::std::optional<at::Tensor> & running_mean, const ::std::optional<at::Tensor> & running_var, double momentum, double eps, const at::Tensor & counts) {
  8284. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8285. auto maybe_layer = maybeCurrentDynamicLayer();
  8286. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8287. int64_t cur_level = maybe_layer->layerId();
  8288. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(mean, cur_level) && !isBatchedAtLevel(invstd, cur_level) && !isBatchedAtLevel(running_mean, cur_level) && !isBatchedAtLevel(running_var, cur_level) && !isBatchedAtLevel(counts, cur_level)) {
  8289. return at::_ops::batch_norm_gather_stats_with_counts::call(input, mean, invstd, running_mean, running_var, momentum, eps, counts);
  8290. }
  8291. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  8292. auto [mean_value, mean_bdim] = unwrapTensorAtLevel(mean, cur_level);
  8293. auto [invstd_value, invstd_bdim] = unwrapTensorAtLevel(invstd, cur_level);
  8294. auto [counts_value, counts_bdim] = unwrapTensorAtLevel(counts, cur_level);
  8295. std::optional<Tensor> running_mean_value;
  8296. std::optional<int64_t> running_mean_bdim;
  8297. if (running_mean) {
  8298. std::tie(running_mean_value, running_mean_bdim) = unwrapTensorAtLevel(running_mean.value(), cur_level);
  8299. }
  8300. std::optional<Tensor> running_var_value;
  8301. std::optional<int64_t> running_var_bdim;
  8302. if (running_var) {
  8303. std::tie(running_var_value, running_var_bdim) = unwrapTensorAtLevel(running_var.value(), cur_level);
  8304. }
  8305. auto results = batch_rule(input_value, input_bdim, mean_value, mean_bdim, invstd_value, invstd_bdim, running_mean_value, running_mean_bdim, running_var_value, running_var_bdim, momentum, eps, counts_value, counts_bdim);
  8306. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  8307. }
  8308. template <typename batch_rule_t, batch_rule_t batch_rule>
  8309. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> native_batch_norm_backward_generated_plumbing(const at::Tensor & grad_out, const at::Tensor & input, const ::std::optional<at::Tensor> & weight, const ::std::optional<at::Tensor> & running_mean, const ::std::optional<at::Tensor> & running_var, const ::std::optional<at::Tensor> & save_mean, const ::std::optional<at::Tensor> & save_invstd, bool train, double eps, ::std::array<bool,3> output_mask) {
  8310. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8311. auto maybe_layer = maybeCurrentDynamicLayer();
  8312. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8313. int64_t cur_level = maybe_layer->layerId();
  8314. if (!isBatchedAtLevel(grad_out, cur_level) && !isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(running_mean, cur_level) && !isBatchedAtLevel(running_var, cur_level) && !isBatchedAtLevel(save_mean, cur_level) && !isBatchedAtLevel(save_invstd, cur_level)) {
  8315. return at::_ops::native_batch_norm_backward::call(grad_out, input, weight, running_mean, running_var, save_mean, save_invstd, train, eps, output_mask);
  8316. }
  8317. auto [grad_out_value, grad_out_bdim] = unwrapTensorAtLevel(grad_out, cur_level);
  8318. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  8319. std::optional<Tensor> weight_value;
  8320. std::optional<int64_t> weight_bdim;
  8321. if (weight) {
  8322. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  8323. }
  8324. std::optional<Tensor> running_mean_value;
  8325. std::optional<int64_t> running_mean_bdim;
  8326. if (running_mean) {
  8327. std::tie(running_mean_value, running_mean_bdim) = unwrapTensorAtLevel(running_mean.value(), cur_level);
  8328. }
  8329. std::optional<Tensor> running_var_value;
  8330. std::optional<int64_t> running_var_bdim;
  8331. if (running_var) {
  8332. std::tie(running_var_value, running_var_bdim) = unwrapTensorAtLevel(running_var.value(), cur_level);
  8333. }
  8334. std::optional<Tensor> save_mean_value;
  8335. std::optional<int64_t> save_mean_bdim;
  8336. if (save_mean) {
  8337. std::tie(save_mean_value, save_mean_bdim) = unwrapTensorAtLevel(save_mean.value(), cur_level);
  8338. }
  8339. std::optional<Tensor> save_invstd_value;
  8340. std::optional<int64_t> save_invstd_bdim;
  8341. if (save_invstd) {
  8342. std::tie(save_invstd_value, save_invstd_bdim) = unwrapTensorAtLevel(save_invstd.value(), cur_level);
  8343. }
  8344. auto results = batch_rule(grad_out_value, grad_out_bdim, input_value, input_bdim, weight_value, weight_bdim, running_mean_value, running_mean_bdim, running_var_value, running_var_bdim, save_mean_value, save_mean_bdim, save_invstd_value, save_invstd_bdim, train, eps, output_mask);
  8345. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  8346. }
  8347. template <typename batch_rule_t, batch_rule_t batch_rule>
  8348. ::std::tuple<at::Tensor,at::Tensor,at::Tensor,at::Tensor> batch_norm_backward_reduce_generated_plumbing(const at::Tensor & grad_out, const at::Tensor & input, const at::Tensor & mean, const at::Tensor & invstd, const ::std::optional<at::Tensor> & weight, bool input_g, bool weight_g, bool bias_g) {
  8349. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8350. auto maybe_layer = maybeCurrentDynamicLayer();
  8351. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8352. int64_t cur_level = maybe_layer->layerId();
  8353. if (!isBatchedAtLevel(grad_out, cur_level) && !isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(mean, cur_level) && !isBatchedAtLevel(invstd, cur_level) && !isBatchedAtLevel(weight, cur_level)) {
  8354. return at::_ops::batch_norm_backward_reduce::call(grad_out, input, mean, invstd, weight, input_g, weight_g, bias_g);
  8355. }
  8356. auto [grad_out_value, grad_out_bdim] = unwrapTensorAtLevel(grad_out, cur_level);
  8357. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  8358. auto [mean_value, mean_bdim] = unwrapTensorAtLevel(mean, cur_level);
  8359. auto [invstd_value, invstd_bdim] = unwrapTensorAtLevel(invstd, cur_level);
  8360. std::optional<Tensor> weight_value;
  8361. std::optional<int64_t> weight_bdim;
  8362. if (weight) {
  8363. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  8364. }
  8365. auto results = batch_rule(grad_out_value, grad_out_bdim, input_value, input_bdim, mean_value, mean_bdim, invstd_value, invstd_bdim, weight_value, weight_bdim, input_g, weight_g, bias_g);
  8366. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level), makeBatched(std::get<6>(results), std::get<7>(results), cur_level));
  8367. }
  8368. template <typename batch_rule_t, batch_rule_t batch_rule>
  8369. at::Tensor batch_norm_backward_elemt_generated_plumbing(const at::Tensor & grad_out, const at::Tensor & input, const at::Tensor & mean, const at::Tensor & invstd, const ::std::optional<at::Tensor> & weight, const at::Tensor & sum_dy, const at::Tensor & sum_dy_xmu, const at::Tensor & count) {
  8370. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8371. auto maybe_layer = maybeCurrentDynamicLayer();
  8372. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8373. int64_t cur_level = maybe_layer->layerId();
  8374. if (!isBatchedAtLevel(grad_out, cur_level) && !isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(mean, cur_level) && !isBatchedAtLevel(invstd, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(sum_dy, cur_level) && !isBatchedAtLevel(sum_dy_xmu, cur_level) && !isBatchedAtLevel(count, cur_level)) {
  8375. return at::_ops::batch_norm_backward_elemt::call(grad_out, input, mean, invstd, weight, sum_dy, sum_dy_xmu, count);
  8376. }
  8377. auto [grad_out_value, grad_out_bdim] = unwrapTensorAtLevel(grad_out, cur_level);
  8378. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  8379. auto [mean_value, mean_bdim] = unwrapTensorAtLevel(mean, cur_level);
  8380. auto [invstd_value, invstd_bdim] = unwrapTensorAtLevel(invstd, cur_level);
  8381. auto [sum_dy_value, sum_dy_bdim] = unwrapTensorAtLevel(sum_dy, cur_level);
  8382. auto [sum_dy_xmu_value, sum_dy_xmu_bdim] = unwrapTensorAtLevel(sum_dy_xmu, cur_level);
  8383. auto [count_value, count_bdim] = unwrapTensorAtLevel(count, cur_level);
  8384. std::optional<Tensor> weight_value;
  8385. std::optional<int64_t> weight_bdim;
  8386. if (weight) {
  8387. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  8388. }
  8389. auto results = batch_rule(grad_out_value, grad_out_bdim, input_value, input_bdim, mean_value, mean_bdim, invstd_value, invstd_bdim, weight_value, weight_bdim, sum_dy_value, sum_dy_bdim, sum_dy_xmu_value, sum_dy_xmu_bdim, count_value, count_bdim);
  8390. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8391. }
  8392. template <typename batch_rule_t, batch_rule_t batch_rule>
  8393. ::std::tuple<at::Tensor,at::Tensor> batch_norm_update_stats_generated_plumbing(const at::Tensor & input, const ::std::optional<at::Tensor> & running_mean, const ::std::optional<at::Tensor> & running_var, double momentum) {
  8394. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8395. auto maybe_layer = maybeCurrentDynamicLayer();
  8396. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8397. int64_t cur_level = maybe_layer->layerId();
  8398. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(running_mean, cur_level) && !isBatchedAtLevel(running_var, cur_level)) {
  8399. return at::_ops::batch_norm_update_stats::call(input, running_mean, running_var, momentum);
  8400. }
  8401. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  8402. std::optional<Tensor> running_mean_value;
  8403. std::optional<int64_t> running_mean_bdim;
  8404. if (running_mean) {
  8405. std::tie(running_mean_value, running_mean_bdim) = unwrapTensorAtLevel(running_mean.value(), cur_level);
  8406. }
  8407. std::optional<Tensor> running_var_value;
  8408. std::optional<int64_t> running_var_bdim;
  8409. if (running_var) {
  8410. std::tie(running_var_value, running_var_bdim) = unwrapTensorAtLevel(running_var.value(), cur_level);
  8411. }
  8412. auto results = batch_rule(input_value, input_bdim, running_mean_value, running_mean_bdim, running_var_value, running_var_bdim, momentum);
  8413. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  8414. }
  8415. template <typename batch_rule_t, batch_rule_t batch_rule>
  8416. at::Tensor _nnpack_spatial_convolution_generated_plumbing(const at::Tensor & input, const at::Tensor & weight, const ::std::optional<at::Tensor> & bias, c10::SymIntArrayRef padding, c10::SymIntArrayRef stride) {
  8417. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8418. auto maybe_layer = maybeCurrentDynamicLayer();
  8419. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8420. int64_t cur_level = maybe_layer->layerId();
  8421. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  8422. return at::_ops::_nnpack_spatial_convolution::call(input, weight, bias, padding, stride);
  8423. }
  8424. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  8425. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  8426. std::optional<Tensor> bias_value;
  8427. std::optional<int64_t> bias_bdim;
  8428. if (bias) {
  8429. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  8430. }
  8431. auto results = batch_rule(input_value, input_bdim, weight_value, weight_bdim, bias_value, bias_bdim, padding, stride);
  8432. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8433. }
  8434. template <typename batch_rule_t, batch_rule_t batch_rule>
  8435. at::Tensor ones_like_generated_plumbing(const at::Tensor & self, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory, ::std::optional<at::MemoryFormat> memory_format) {
  8436. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8437. auto maybe_layer = maybeCurrentDynamicLayer();
  8438. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8439. int64_t cur_level = maybe_layer->layerId();
  8440. if (!isBatchedAtLevel(self, cur_level)) {
  8441. return at::_ops::ones_like::call(self, dtype, layout, device, pin_memory, memory_format);
  8442. }
  8443. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8444. auto results = batch_rule(self_value, self_bdim, dtype, layout, device, pin_memory, memory_format);
  8445. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8446. }
  8447. template <typename batch_rule_t, batch_rule_t batch_rule>
  8448. at::Tensor pairwise_distance_generated_plumbing(const at::Tensor & x1, const at::Tensor & x2, double p, double eps, bool keepdim) {
  8449. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8450. auto maybe_layer = maybeCurrentDynamicLayer();
  8451. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8452. int64_t cur_level = maybe_layer->layerId();
  8453. if (!isBatchedAtLevel(x1, cur_level) && !isBatchedAtLevel(x2, cur_level)) {
  8454. return at::_ops::pairwise_distance::call(x1, x2, p, eps, keepdim);
  8455. }
  8456. auto [x1_value, x1_bdim] = unwrapTensorAtLevel(x1, cur_level);
  8457. auto [x2_value, x2_bdim] = unwrapTensorAtLevel(x2, cur_level);
  8458. auto results = batch_rule(x1_value, x1_bdim, x2_value, x2_bdim, p, eps, keepdim);
  8459. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8460. }
  8461. template <typename batch_rule_t, batch_rule_t batch_rule>
  8462. at::Tensor cdist_generated_plumbing(const at::Tensor & x1, const at::Tensor & x2, double p, ::std::optional<int64_t> compute_mode) {
  8463. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8464. auto maybe_layer = maybeCurrentDynamicLayer();
  8465. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8466. int64_t cur_level = maybe_layer->layerId();
  8467. if (!isBatchedAtLevel(x1, cur_level) && !isBatchedAtLevel(x2, cur_level)) {
  8468. return at::_ops::cdist::call(x1, x2, p, compute_mode);
  8469. }
  8470. auto [x1_value, x1_bdim] = unwrapTensorAtLevel(x1, cur_level);
  8471. auto [x2_value, x2_bdim] = unwrapTensorAtLevel(x2, cur_level);
  8472. auto results = batch_rule(x1_value, x1_bdim, x2_value, x2_bdim, p, compute_mode);
  8473. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8474. }
  8475. template <typename batch_rule_t, batch_rule_t batch_rule>
  8476. at::Tensor _euclidean_dist_generated_plumbing(const at::Tensor & x1, const at::Tensor & x2) {
  8477. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8478. auto maybe_layer = maybeCurrentDynamicLayer();
  8479. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8480. int64_t cur_level = maybe_layer->layerId();
  8481. if (!isBatchedAtLevel(x1, cur_level) && !isBatchedAtLevel(x2, cur_level)) {
  8482. return at::_ops::_euclidean_dist::call(x1, x2);
  8483. }
  8484. auto [x1_value, x1_bdim] = unwrapTensorAtLevel(x1, cur_level);
  8485. auto [x2_value, x2_bdim] = unwrapTensorAtLevel(x2, cur_level);
  8486. auto results = batch_rule(x1_value, x1_bdim, x2_value, x2_bdim);
  8487. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8488. }
  8489. template <typename batch_rule_t, batch_rule_t batch_rule>
  8490. at::Tensor _cdist_forward_generated_plumbing(const at::Tensor & x1, const at::Tensor & x2, double p, ::std::optional<int64_t> compute_mode) {
  8491. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8492. auto maybe_layer = maybeCurrentDynamicLayer();
  8493. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8494. int64_t cur_level = maybe_layer->layerId();
  8495. if (!isBatchedAtLevel(x1, cur_level) && !isBatchedAtLevel(x2, cur_level)) {
  8496. return at::_ops::_cdist_forward::call(x1, x2, p, compute_mode);
  8497. }
  8498. auto [x1_value, x1_bdim] = unwrapTensorAtLevel(x1, cur_level);
  8499. auto [x2_value, x2_bdim] = unwrapTensorAtLevel(x2, cur_level);
  8500. auto results = batch_rule(x1_value, x1_bdim, x2_value, x2_bdim, p, compute_mode);
  8501. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8502. }
  8503. template <typename batch_rule_t, batch_rule_t batch_rule>
  8504. at::Tensor _cdist_backward_generated_plumbing(const at::Tensor & grad, const at::Tensor & x1, const at::Tensor & x2, double p, const at::Tensor & cdist) {
  8505. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8506. auto maybe_layer = maybeCurrentDynamicLayer();
  8507. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8508. int64_t cur_level = maybe_layer->layerId();
  8509. if (!isBatchedAtLevel(grad, cur_level) && !isBatchedAtLevel(x1, cur_level) && !isBatchedAtLevel(x2, cur_level) && !isBatchedAtLevel(cdist, cur_level)) {
  8510. return at::_ops::_cdist_backward::call(grad, x1, x2, p, cdist);
  8511. }
  8512. auto [grad_value, grad_bdim] = unwrapTensorAtLevel(grad, cur_level);
  8513. auto [x1_value, x1_bdim] = unwrapTensorAtLevel(x1, cur_level);
  8514. auto [x2_value, x2_bdim] = unwrapTensorAtLevel(x2, cur_level);
  8515. auto [cdist_value, cdist_bdim] = unwrapTensorAtLevel(cdist, cur_level);
  8516. auto results = batch_rule(grad_value, grad_bdim, x1_value, x1_bdim, x2_value, x2_bdim, p, cdist_value, cdist_bdim);
  8517. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8518. }
  8519. template <typename batch_rule_t, batch_rule_t batch_rule>
  8520. at::Tensor pdist_generated_plumbing(const at::Tensor & self, double p) {
  8521. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8522. auto maybe_layer = maybeCurrentDynamicLayer();
  8523. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8524. int64_t cur_level = maybe_layer->layerId();
  8525. if (!isBatchedAtLevel(self, cur_level)) {
  8526. return at::_ops::pdist::call(self, p);
  8527. }
  8528. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8529. auto results = batch_rule(self_value, self_bdim, p);
  8530. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8531. }
  8532. template <typename batch_rule_t, batch_rule_t batch_rule>
  8533. at::Tensor _pdist_forward_generated_plumbing(const at::Tensor & self, double p) {
  8534. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8535. auto maybe_layer = maybeCurrentDynamicLayer();
  8536. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8537. int64_t cur_level = maybe_layer->layerId();
  8538. if (!isBatchedAtLevel(self, cur_level)) {
  8539. return at::_ops::_pdist_forward::call(self, p);
  8540. }
  8541. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8542. auto results = batch_rule(self_value, self_bdim, p);
  8543. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8544. }
  8545. template <typename batch_rule_t, batch_rule_t batch_rule>
  8546. at::Tensor _pdist_backward_generated_plumbing(const at::Tensor & grad, const at::Tensor & self, double p, const at::Tensor & pdist) {
  8547. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8548. auto maybe_layer = maybeCurrentDynamicLayer();
  8549. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8550. int64_t cur_level = maybe_layer->layerId();
  8551. if (!isBatchedAtLevel(grad, cur_level) && !isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(pdist, cur_level)) {
  8552. return at::_ops::_pdist_backward::call(grad, self, p, pdist);
  8553. }
  8554. auto [grad_value, grad_bdim] = unwrapTensorAtLevel(grad, cur_level);
  8555. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8556. auto [pdist_value, pdist_bdim] = unwrapTensorAtLevel(pdist, cur_level);
  8557. auto results = batch_rule(grad_value, grad_bdim, self_value, self_bdim, p, pdist_value, pdist_bdim);
  8558. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8559. }
  8560. template <typename batch_rule_t, batch_rule_t batch_rule>
  8561. at::Tensor cosine_similarity_generated_plumbing(const at::Tensor & x1, const at::Tensor & x2, int64_t dim, double eps) {
  8562. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8563. auto maybe_layer = maybeCurrentDynamicLayer();
  8564. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8565. int64_t cur_level = maybe_layer->layerId();
  8566. if (!isBatchedAtLevel(x1, cur_level) && !isBatchedAtLevel(x2, cur_level)) {
  8567. return at::_ops::cosine_similarity::call(x1, x2, dim, eps);
  8568. }
  8569. auto [x1_value, x1_bdim] = unwrapTensorAtLevel(x1, cur_level);
  8570. auto [x2_value, x2_bdim] = unwrapTensorAtLevel(x2, cur_level);
  8571. auto results = batch_rule(x1_value, x1_bdim, x2_value, x2_bdim, dim, eps);
  8572. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8573. }
  8574. template <typename batch_rule_t, batch_rule_t batch_rule>
  8575. at::Tensor permute_generated_plumbing(const at::Tensor & self, at::IntArrayRef dims) {
  8576. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8577. auto maybe_layer = maybeCurrentDynamicLayer();
  8578. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8579. int64_t cur_level = maybe_layer->layerId();
  8580. if (!isBatchedAtLevel(self, cur_level)) {
  8581. return at::_ops::permute::call(self, dims);
  8582. }
  8583. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8584. auto results = batch_rule(self_value, self_bdim, dims);
  8585. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8586. }
  8587. template <typename batch_rule_t, batch_rule_t batch_rule>
  8588. at::Tensor movedim_intlist_generated_plumbing(const at::Tensor & self, at::IntArrayRef source, at::IntArrayRef destination) {
  8589. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8590. auto maybe_layer = maybeCurrentDynamicLayer();
  8591. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8592. int64_t cur_level = maybe_layer->layerId();
  8593. if (!isBatchedAtLevel(self, cur_level)) {
  8594. return at::_ops::movedim_intlist::call(self, source, destination);
  8595. }
  8596. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8597. auto results = batch_rule(self_value, self_bdim, source, destination);
  8598. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8599. }
  8600. template <typename batch_rule_t, batch_rule_t batch_rule>
  8601. at::Tensor movedim_int_generated_plumbing(const at::Tensor & self, int64_t source, int64_t destination) {
  8602. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8603. auto maybe_layer = maybeCurrentDynamicLayer();
  8604. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8605. int64_t cur_level = maybe_layer->layerId();
  8606. if (!isBatchedAtLevel(self, cur_level)) {
  8607. return at::_ops::movedim_int::call(self, source, destination);
  8608. }
  8609. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8610. auto results = batch_rule(self_value, self_bdim, source, destination);
  8611. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8612. }
  8613. template <typename batch_rule_t, batch_rule_t batch_rule>
  8614. at::Tensor moveaxis_intlist_generated_plumbing(const at::Tensor & self, at::IntArrayRef source, at::IntArrayRef destination) {
  8615. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8616. auto maybe_layer = maybeCurrentDynamicLayer();
  8617. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8618. int64_t cur_level = maybe_layer->layerId();
  8619. if (!isBatchedAtLevel(self, cur_level)) {
  8620. return at::_ops::moveaxis_intlist::call(self, source, destination);
  8621. }
  8622. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8623. auto results = batch_rule(self_value, self_bdim, source, destination);
  8624. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8625. }
  8626. template <typename batch_rule_t, batch_rule_t batch_rule>
  8627. at::Tensor moveaxis_int_generated_plumbing(const at::Tensor & self, int64_t source, int64_t destination) {
  8628. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8629. auto maybe_layer = maybeCurrentDynamicLayer();
  8630. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8631. int64_t cur_level = maybe_layer->layerId();
  8632. if (!isBatchedAtLevel(self, cur_level)) {
  8633. return at::_ops::moveaxis_int::call(self, source, destination);
  8634. }
  8635. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8636. auto results = batch_rule(self_value, self_bdim, source, destination);
  8637. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8638. }
  8639. template <typename batch_rule_t, batch_rule_t batch_rule>
  8640. at::Tensor numpy_T_generated_plumbing(const at::Tensor & self) {
  8641. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8642. auto maybe_layer = maybeCurrentDynamicLayer();
  8643. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8644. int64_t cur_level = maybe_layer->layerId();
  8645. if (!isBatchedAtLevel(self, cur_level)) {
  8646. return at::_ops::numpy_T::call(self);
  8647. }
  8648. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8649. auto results = batch_rule(self_value, self_bdim);
  8650. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8651. }
  8652. template <typename batch_rule_t, batch_rule_t batch_rule>
  8653. at::Tensor matrix_H_generated_plumbing(const at::Tensor & self) {
  8654. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8655. auto maybe_layer = maybeCurrentDynamicLayer();
  8656. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8657. int64_t cur_level = maybe_layer->layerId();
  8658. if (!isBatchedAtLevel(self, cur_level)) {
  8659. return at::_ops::matrix_H::call(self);
  8660. }
  8661. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8662. auto results = batch_rule(self_value, self_bdim);
  8663. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8664. }
  8665. template <typename batch_rule_t, batch_rule_t batch_rule>
  8666. at::Tensor mT_generated_plumbing(const at::Tensor & self) {
  8667. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8668. auto maybe_layer = maybeCurrentDynamicLayer();
  8669. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8670. int64_t cur_level = maybe_layer->layerId();
  8671. if (!isBatchedAtLevel(self, cur_level)) {
  8672. return at::_ops::mT::call(self);
  8673. }
  8674. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8675. auto results = batch_rule(self_value, self_bdim);
  8676. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8677. }
  8678. template <typename batch_rule_t, batch_rule_t batch_rule>
  8679. at::Tensor mH_generated_plumbing(const at::Tensor & self) {
  8680. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8681. auto maybe_layer = maybeCurrentDynamicLayer();
  8682. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8683. int64_t cur_level = maybe_layer->layerId();
  8684. if (!isBatchedAtLevel(self, cur_level)) {
  8685. return at::_ops::mH::call(self);
  8686. }
  8687. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8688. auto results = batch_rule(self_value, self_bdim);
  8689. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8690. }
  8691. template <typename batch_rule_t, batch_rule_t batch_rule>
  8692. at::Tensor adjoint_generated_plumbing(const at::Tensor & self) {
  8693. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8694. auto maybe_layer = maybeCurrentDynamicLayer();
  8695. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8696. int64_t cur_level = maybe_layer->layerId();
  8697. if (!isBatchedAtLevel(self, cur_level)) {
  8698. return at::_ops::adjoint::call(self);
  8699. }
  8700. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8701. auto results = batch_rule(self_value, self_bdim);
  8702. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8703. }
  8704. template <typename batch_rule_t, batch_rule_t batch_rule>
  8705. at::Tensor pixel_shuffle_generated_plumbing(const at::Tensor & self, int64_t upscale_factor) {
  8706. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8707. auto maybe_layer = maybeCurrentDynamicLayer();
  8708. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8709. int64_t cur_level = maybe_layer->layerId();
  8710. if (!isBatchedAtLevel(self, cur_level)) {
  8711. return at::_ops::pixel_shuffle::call(self, upscale_factor);
  8712. }
  8713. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8714. auto results = batch_rule(self_value, self_bdim, upscale_factor);
  8715. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8716. }
  8717. template <typename batch_rule_t, batch_rule_t batch_rule>
  8718. at::Tensor pixel_unshuffle_generated_plumbing(const at::Tensor & self, int64_t downscale_factor) {
  8719. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8720. auto maybe_layer = maybeCurrentDynamicLayer();
  8721. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8722. int64_t cur_level = maybe_layer->layerId();
  8723. if (!isBatchedAtLevel(self, cur_level)) {
  8724. return at::_ops::pixel_unshuffle::call(self, downscale_factor);
  8725. }
  8726. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8727. auto results = batch_rule(self_value, self_bdim, downscale_factor);
  8728. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8729. }
  8730. template <typename batch_rule_t, batch_rule_t batch_rule>
  8731. at::Tensor channel_shuffle_generated_plumbing(const at::Tensor & self, c10::SymInt groups) {
  8732. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8733. auto maybe_layer = maybeCurrentDynamicLayer();
  8734. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8735. int64_t cur_level = maybe_layer->layerId();
  8736. if (!isBatchedAtLevel(self, cur_level)) {
  8737. return at::_ops::channel_shuffle::call(self, groups);
  8738. }
  8739. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8740. auto results = batch_rule(self_value, self_bdim, groups);
  8741. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8742. }
  8743. template <typename batch_rule_t, batch_rule_t batch_rule>
  8744. at::Tensor native_channel_shuffle_generated_plumbing(const at::Tensor & self, c10::SymInt groups) {
  8745. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8746. auto maybe_layer = maybeCurrentDynamicLayer();
  8747. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8748. int64_t cur_level = maybe_layer->layerId();
  8749. if (!isBatchedAtLevel(self, cur_level)) {
  8750. return at::_ops::native_channel_shuffle::call(self, groups);
  8751. }
  8752. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8753. auto results = batch_rule(self_value, self_bdim, groups);
  8754. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8755. }
  8756. template <typename batch_rule_t, batch_rule_t batch_rule>
  8757. at::Tensor pin_memory_generated_plumbing(const at::Tensor & self, ::std::optional<at::Device> device) {
  8758. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8759. auto maybe_layer = maybeCurrentDynamicLayer();
  8760. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8761. int64_t cur_level = maybe_layer->layerId();
  8762. if (!isBatchedAtLevel(self, cur_level)) {
  8763. return at::_ops::pin_memory::call(self, device);
  8764. }
  8765. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8766. auto results = batch_rule(self_value, self_bdim, device);
  8767. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8768. }
  8769. template <typename batch_rule_t, batch_rule_t batch_rule>
  8770. at::Tensor _pin_memory_generated_plumbing(const at::Tensor & self, ::std::optional<at::Device> device) {
  8771. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8772. auto maybe_layer = maybeCurrentDynamicLayer();
  8773. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8774. int64_t cur_level = maybe_layer->layerId();
  8775. if (!isBatchedAtLevel(self, cur_level)) {
  8776. return at::_ops::_pin_memory::call(self, device);
  8777. }
  8778. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8779. auto results = batch_rule(self_value, self_bdim, device);
  8780. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8781. }
  8782. template <typename batch_rule_t, batch_rule_t batch_rule>
  8783. at::Tensor pinverse_generated_plumbing(const at::Tensor & self, double rcond) {
  8784. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8785. auto maybe_layer = maybeCurrentDynamicLayer();
  8786. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8787. int64_t cur_level = maybe_layer->layerId();
  8788. if (!isBatchedAtLevel(self, cur_level)) {
  8789. return at::_ops::pinverse::call(self, rcond);
  8790. }
  8791. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8792. auto results = batch_rule(self_value, self_bdim, rcond);
  8793. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8794. }
  8795. template <typename batch_rule_t, batch_rule_t batch_rule>
  8796. at::Tensor poisson_nll_loss_generated_plumbing(const at::Tensor & input, const at::Tensor & target, bool log_input, bool full, double eps, int64_t reduction) {
  8797. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8798. auto maybe_layer = maybeCurrentDynamicLayer();
  8799. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8800. int64_t cur_level = maybe_layer->layerId();
  8801. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(target, cur_level)) {
  8802. return at::_ops::poisson_nll_loss::call(input, target, log_input, full, eps, reduction);
  8803. }
  8804. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  8805. auto [target_value, target_bdim] = unwrapTensorAtLevel(target, cur_level);
  8806. auto results = batch_rule(input_value, input_bdim, target_value, target_bdim, log_input, full, eps, reduction);
  8807. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8808. }
  8809. template <typename batch_rule_t, batch_rule_t batch_rule>
  8810. at::Tensor rad2deg_generated_plumbing(const at::Tensor & self) {
  8811. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8812. auto maybe_layer = maybeCurrentDynamicLayer();
  8813. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8814. int64_t cur_level = maybe_layer->layerId();
  8815. if (!isBatchedAtLevel(self, cur_level)) {
  8816. return at::_ops::rad2deg::call(self);
  8817. }
  8818. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8819. auto results = batch_rule(self_value, self_bdim);
  8820. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8821. }
  8822. template <typename batch_rule_t, batch_rule_t batch_rule>
  8823. at::Tensor & rad2deg__generated_plumbing(at::Tensor & self) {
  8824. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8825. auto maybe_layer = maybeCurrentDynamicLayer();
  8826. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  8827. int64_t cur_level = maybe_layer->layerId();
  8828. if (!isBatchedAtLevel(self, cur_level)) {
  8829. return at::_ops::rad2deg_::call(self);
  8830. }
  8831. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8832. batch_rule(self_value, self_bdim);
  8833. return self;
  8834. }
  8835. template <typename batch_rule_t, batch_rule_t batch_rule>
  8836. at::Tensor deg2rad_generated_plumbing(const at::Tensor & self) {
  8837. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8838. auto maybe_layer = maybeCurrentDynamicLayer();
  8839. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8840. int64_t cur_level = maybe_layer->layerId();
  8841. if (!isBatchedAtLevel(self, cur_level)) {
  8842. return at::_ops::deg2rad::call(self);
  8843. }
  8844. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8845. auto results = batch_rule(self_value, self_bdim);
  8846. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8847. }
  8848. template <typename batch_rule_t, batch_rule_t batch_rule>
  8849. at::Tensor & deg2rad__generated_plumbing(at::Tensor & self) {
  8850. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8851. auto maybe_layer = maybeCurrentDynamicLayer();
  8852. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  8853. int64_t cur_level = maybe_layer->layerId();
  8854. if (!isBatchedAtLevel(self, cur_level)) {
  8855. return at::_ops::deg2rad_::call(self);
  8856. }
  8857. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8858. batch_rule(self_value, self_bdim);
  8859. return self;
  8860. }
  8861. template <typename batch_rule_t, batch_rule_t batch_rule>
  8862. at::Tensor rand_like_generated_plumbing(const at::Tensor & self, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory, ::std::optional<at::MemoryFormat> memory_format) {
  8863. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8864. auto maybe_layer = maybeCurrentDynamicLayer();
  8865. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8866. int64_t cur_level = maybe_layer->layerId();
  8867. if (!isBatchedAtLevel(self, cur_level)) {
  8868. return at::_ops::rand_like::call(self, dtype, layout, device, pin_memory, memory_format);
  8869. }
  8870. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8871. auto results = batch_rule(self_value, self_bdim, dtype, layout, device, pin_memory, memory_format);
  8872. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8873. }
  8874. template <typename batch_rule_t, batch_rule_t batch_rule>
  8875. at::Tensor randint_like_generated_plumbing(const at::Tensor & self, c10::SymInt high, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory, ::std::optional<at::MemoryFormat> memory_format) {
  8876. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8877. auto maybe_layer = maybeCurrentDynamicLayer();
  8878. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8879. int64_t cur_level = maybe_layer->layerId();
  8880. if (!isBatchedAtLevel(self, cur_level)) {
  8881. return at::_ops::randint_like::call(self, high, dtype, layout, device, pin_memory, memory_format);
  8882. }
  8883. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8884. auto results = batch_rule(self_value, self_bdim, high, dtype, layout, device, pin_memory, memory_format);
  8885. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8886. }
  8887. template <typename batch_rule_t, batch_rule_t batch_rule>
  8888. at::Tensor randint_like_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & high, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory, ::std::optional<at::MemoryFormat> memory_format) {
  8889. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8890. auto maybe_layer = maybeCurrentDynamicLayer();
  8891. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8892. int64_t cur_level = maybe_layer->layerId();
  8893. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(high, cur_level)) {
  8894. return at::_ops::randint_like_Tensor::call(self, high, dtype, layout, device, pin_memory, memory_format);
  8895. }
  8896. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8897. auto [high_value, high_bdim] = unwrapTensorAtLevel(high, cur_level);
  8898. auto results = batch_rule(self_value, self_bdim, high_value, high_bdim, dtype, layout, device, pin_memory, memory_format);
  8899. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8900. }
  8901. template <typename batch_rule_t, batch_rule_t batch_rule>
  8902. at::Tensor randint_like_low_dtype_generated_plumbing(const at::Tensor & self, c10::SymInt low, c10::SymInt high, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory, ::std::optional<at::MemoryFormat> memory_format) {
  8903. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8904. auto maybe_layer = maybeCurrentDynamicLayer();
  8905. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8906. int64_t cur_level = maybe_layer->layerId();
  8907. if (!isBatchedAtLevel(self, cur_level)) {
  8908. return at::_ops::randint_like_low_dtype::call(self, low, high, dtype, layout, device, pin_memory, memory_format);
  8909. }
  8910. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8911. auto results = batch_rule(self_value, self_bdim, low, high, dtype, layout, device, pin_memory, memory_format);
  8912. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8913. }
  8914. template <typename batch_rule_t, batch_rule_t batch_rule>
  8915. at::Tensor randn_like_generated_plumbing(const at::Tensor & self, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory, ::std::optional<at::MemoryFormat> memory_format) {
  8916. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8917. auto maybe_layer = maybeCurrentDynamicLayer();
  8918. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8919. int64_t cur_level = maybe_layer->layerId();
  8920. if (!isBatchedAtLevel(self, cur_level)) {
  8921. return at::_ops::randn_like::call(self, dtype, layout, device, pin_memory, memory_format);
  8922. }
  8923. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8924. auto results = batch_rule(self_value, self_bdim, dtype, layout, device, pin_memory, memory_format);
  8925. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8926. }
  8927. template <typename batch_rule_t, batch_rule_t batch_rule>
  8928. at::Tensor ravel_generated_plumbing(const at::Tensor & self) {
  8929. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8930. auto maybe_layer = maybeCurrentDynamicLayer();
  8931. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8932. int64_t cur_level = maybe_layer->layerId();
  8933. if (!isBatchedAtLevel(self, cur_level)) {
  8934. return at::_ops::ravel::call(self);
  8935. }
  8936. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8937. auto results = batch_rule(self_value, self_bdim);
  8938. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8939. }
  8940. template <typename batch_rule_t, batch_rule_t batch_rule>
  8941. at::Tensor reciprocal_generated_plumbing(const at::Tensor & self) {
  8942. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8943. auto maybe_layer = maybeCurrentDynamicLayer();
  8944. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8945. int64_t cur_level = maybe_layer->layerId();
  8946. if (!isBatchedAtLevel(self, cur_level)) {
  8947. return at::_ops::reciprocal::call(self);
  8948. }
  8949. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8950. auto results = batch_rule(self_value, self_bdim);
  8951. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8952. }
  8953. template <typename batch_rule_t, batch_rule_t batch_rule>
  8954. at::Tensor & reciprocal__generated_plumbing(at::Tensor & self) {
  8955. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8956. auto maybe_layer = maybeCurrentDynamicLayer();
  8957. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  8958. int64_t cur_level = maybe_layer->layerId();
  8959. if (!isBatchedAtLevel(self, cur_level)) {
  8960. return at::_ops::reciprocal_::call(self);
  8961. }
  8962. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8963. batch_rule(self_value, self_bdim);
  8964. return self;
  8965. }
  8966. template <typename batch_rule_t, batch_rule_t batch_rule>
  8967. at::Tensor neg_generated_plumbing(const at::Tensor & self) {
  8968. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8969. auto maybe_layer = maybeCurrentDynamicLayer();
  8970. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8971. int64_t cur_level = maybe_layer->layerId();
  8972. if (!isBatchedAtLevel(self, cur_level)) {
  8973. return at::_ops::neg::call(self);
  8974. }
  8975. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8976. auto results = batch_rule(self_value, self_bdim);
  8977. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  8978. }
  8979. template <typename batch_rule_t, batch_rule_t batch_rule>
  8980. at::Tensor & neg__generated_plumbing(at::Tensor & self) {
  8981. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8982. auto maybe_layer = maybeCurrentDynamicLayer();
  8983. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  8984. int64_t cur_level = maybe_layer->layerId();
  8985. if (!isBatchedAtLevel(self, cur_level)) {
  8986. return at::_ops::neg_::call(self);
  8987. }
  8988. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  8989. batch_rule(self_value, self_bdim);
  8990. return self;
  8991. }
  8992. template <typename batch_rule_t, batch_rule_t batch_rule>
  8993. at::Tensor negative_generated_plumbing(const at::Tensor & self) {
  8994. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  8995. auto maybe_layer = maybeCurrentDynamicLayer();
  8996. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  8997. int64_t cur_level = maybe_layer->layerId();
  8998. if (!isBatchedAtLevel(self, cur_level)) {
  8999. return at::_ops::negative::call(self);
  9000. }
  9001. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9002. auto results = batch_rule(self_value, self_bdim);
  9003. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9004. }
  9005. template <typename batch_rule_t, batch_rule_t batch_rule>
  9006. at::Tensor & negative__generated_plumbing(at::Tensor & self) {
  9007. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9008. auto maybe_layer = maybeCurrentDynamicLayer();
  9009. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  9010. int64_t cur_level = maybe_layer->layerId();
  9011. if (!isBatchedAtLevel(self, cur_level)) {
  9012. return at::_ops::negative_::call(self);
  9013. }
  9014. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9015. batch_rule(self_value, self_bdim);
  9016. return self;
  9017. }
  9018. template <typename batch_rule_t, batch_rule_t batch_rule>
  9019. at::Tensor repeat_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef repeats) {
  9020. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9021. auto maybe_layer = maybeCurrentDynamicLayer();
  9022. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9023. int64_t cur_level = maybe_layer->layerId();
  9024. if (!isBatchedAtLevel(self, cur_level)) {
  9025. return at::_ops::repeat::call(self, repeats);
  9026. }
  9027. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9028. auto results = batch_rule(self_value, self_bdim, repeats);
  9029. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9030. }
  9031. template <typename batch_rule_t, batch_rule_t batch_rule>
  9032. at::Tensor repeat_interleave_Tensor_generated_plumbing(const at::Tensor & repeats, ::std::optional<c10::SymInt> output_size) {
  9033. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9034. auto maybe_layer = maybeCurrentDynamicLayer();
  9035. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9036. int64_t cur_level = maybe_layer->layerId();
  9037. if (!isBatchedAtLevel(repeats, cur_level)) {
  9038. return at::_ops::repeat_interleave_Tensor::call(repeats, output_size);
  9039. }
  9040. auto [repeats_value, repeats_bdim] = unwrapTensorAtLevel(repeats, cur_level);
  9041. auto results = batch_rule(repeats_value, repeats_bdim, output_size);
  9042. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9043. }
  9044. template <typename batch_rule_t, batch_rule_t batch_rule>
  9045. at::Tensor repeat_interleave_self_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & repeats, ::std::optional<int64_t> dim, ::std::optional<c10::SymInt> output_size) {
  9046. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9047. auto maybe_layer = maybeCurrentDynamicLayer();
  9048. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9049. int64_t cur_level = maybe_layer->layerId();
  9050. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(repeats, cur_level)) {
  9051. return at::_ops::repeat_interleave_self_Tensor::call(self, repeats, dim, output_size);
  9052. }
  9053. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9054. auto [repeats_value, repeats_bdim] = unwrapTensorAtLevel(repeats, cur_level);
  9055. auto results = batch_rule(self_value, self_bdim, repeats_value, repeats_bdim, dim, output_size);
  9056. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9057. }
  9058. template <typename batch_rule_t, batch_rule_t batch_rule>
  9059. at::Tensor repeat_interleave_self_int_generated_plumbing(const at::Tensor & self, c10::SymInt repeats, ::std::optional<int64_t> dim, ::std::optional<c10::SymInt> output_size) {
  9060. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9061. auto maybe_layer = maybeCurrentDynamicLayer();
  9062. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9063. int64_t cur_level = maybe_layer->layerId();
  9064. if (!isBatchedAtLevel(self, cur_level)) {
  9065. return at::_ops::repeat_interleave_self_int::call(self, repeats, dim, output_size);
  9066. }
  9067. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9068. auto results = batch_rule(self_value, self_bdim, repeats, dim, output_size);
  9069. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9070. }
  9071. template <typename batch_rule_t, batch_rule_t batch_rule>
  9072. at::Tensor reshape_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef shape) {
  9073. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9074. auto maybe_layer = maybeCurrentDynamicLayer();
  9075. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9076. int64_t cur_level = maybe_layer->layerId();
  9077. if (!isBatchedAtLevel(self, cur_level)) {
  9078. return at::_ops::reshape::call(self, shape);
  9079. }
  9080. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9081. auto results = batch_rule(self_value, self_bdim, shape);
  9082. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9083. }
  9084. template <typename batch_rule_t, batch_rule_t batch_rule>
  9085. at::Tensor _reshape_copy_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef size) {
  9086. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9087. auto maybe_layer = maybeCurrentDynamicLayer();
  9088. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9089. int64_t cur_level = maybe_layer->layerId();
  9090. if (!isBatchedAtLevel(self, cur_level)) {
  9091. return at::_ops::_reshape_copy::call(self, size);
  9092. }
  9093. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9094. auto results = batch_rule(self_value, self_bdim, size);
  9095. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9096. }
  9097. template <typename batch_rule_t, batch_rule_t batch_rule>
  9098. at::Tensor _reshape_alias_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef size, c10::SymIntArrayRef stride) {
  9099. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9100. auto maybe_layer = maybeCurrentDynamicLayer();
  9101. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9102. int64_t cur_level = maybe_layer->layerId();
  9103. if (!isBatchedAtLevel(self, cur_level)) {
  9104. return at::_ops::_reshape_alias::call(self, size, stride);
  9105. }
  9106. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9107. auto results = batch_rule(self_value, self_bdim, size, stride);
  9108. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9109. }
  9110. template <typename batch_rule_t, batch_rule_t batch_rule>
  9111. at::Tensor _mkldnn_reshape_generated_plumbing(const at::Tensor & self, at::IntArrayRef shape) {
  9112. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9113. auto maybe_layer = maybeCurrentDynamicLayer();
  9114. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9115. int64_t cur_level = maybe_layer->layerId();
  9116. if (!isBatchedAtLevel(self, cur_level)) {
  9117. return at::_ops::_mkldnn_reshape::call(self, shape);
  9118. }
  9119. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9120. auto results = batch_rule(self_value, self_bdim, shape);
  9121. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9122. }
  9123. template <typename batch_rule_t, batch_rule_t batch_rule>
  9124. at::Tensor reshape_as_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  9125. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9126. auto maybe_layer = maybeCurrentDynamicLayer();
  9127. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9128. int64_t cur_level = maybe_layer->layerId();
  9129. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  9130. return at::_ops::reshape_as::call(self, other);
  9131. }
  9132. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9133. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  9134. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  9135. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9136. }
  9137. template <typename batch_rule_t, batch_rule_t batch_rule>
  9138. at::Tensor round_generated_plumbing(const at::Tensor & self) {
  9139. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9140. auto maybe_layer = maybeCurrentDynamicLayer();
  9141. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9142. int64_t cur_level = maybe_layer->layerId();
  9143. if (!isBatchedAtLevel(self, cur_level)) {
  9144. return at::_ops::round::call(self);
  9145. }
  9146. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9147. auto results = batch_rule(self_value, self_bdim);
  9148. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9149. }
  9150. template <typename batch_rule_t, batch_rule_t batch_rule>
  9151. at::Tensor & round__generated_plumbing(at::Tensor & self) {
  9152. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9153. auto maybe_layer = maybeCurrentDynamicLayer();
  9154. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  9155. int64_t cur_level = maybe_layer->layerId();
  9156. if (!isBatchedAtLevel(self, cur_level)) {
  9157. return at::_ops::round_::call(self);
  9158. }
  9159. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9160. batch_rule(self_value, self_bdim);
  9161. return self;
  9162. }
  9163. template <typename batch_rule_t, batch_rule_t batch_rule>
  9164. at::Tensor round_decimals_generated_plumbing(const at::Tensor & self, int64_t decimals) {
  9165. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9166. auto maybe_layer = maybeCurrentDynamicLayer();
  9167. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9168. int64_t cur_level = maybe_layer->layerId();
  9169. if (!isBatchedAtLevel(self, cur_level)) {
  9170. return at::_ops::round_decimals::call(self, decimals);
  9171. }
  9172. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9173. auto results = batch_rule(self_value, self_bdim, decimals);
  9174. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9175. }
  9176. template <typename batch_rule_t, batch_rule_t batch_rule>
  9177. at::Tensor & round__decimals_generated_plumbing(at::Tensor & self, int64_t decimals) {
  9178. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9179. auto maybe_layer = maybeCurrentDynamicLayer();
  9180. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  9181. int64_t cur_level = maybe_layer->layerId();
  9182. if (!isBatchedAtLevel(self, cur_level)) {
  9183. return at::_ops::round__decimals::call(self, decimals);
  9184. }
  9185. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9186. batch_rule(self_value, self_bdim, decimals);
  9187. return self;
  9188. }
  9189. template <typename batch_rule_t, batch_rule_t batch_rule>
  9190. at::Tensor rrelu_generated_plumbing(const at::Tensor & self, const at::Scalar & lower, const at::Scalar & upper, bool training, ::std::optional<at::Generator> generator) {
  9191. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9192. auto maybe_layer = maybeCurrentDynamicLayer();
  9193. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9194. int64_t cur_level = maybe_layer->layerId();
  9195. if (!isBatchedAtLevel(self, cur_level)) {
  9196. return at::_ops::rrelu::call(self, lower, upper, training, generator);
  9197. }
  9198. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9199. auto results = batch_rule(self_value, self_bdim, lower, upper, training, generator);
  9200. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9201. }
  9202. template <typename batch_rule_t, batch_rule_t batch_rule>
  9203. at::Tensor & rrelu__generated_plumbing(at::Tensor & self, const at::Scalar & lower, const at::Scalar & upper, bool training, ::std::optional<at::Generator> generator) {
  9204. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9205. auto maybe_layer = maybeCurrentDynamicLayer();
  9206. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  9207. int64_t cur_level = maybe_layer->layerId();
  9208. if (!isBatchedAtLevel(self, cur_level)) {
  9209. return at::_ops::rrelu_::call(self, lower, upper, training, generator);
  9210. }
  9211. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9212. batch_rule(self_value, self_bdim, lower, upper, training, generator);
  9213. return self;
  9214. }
  9215. template <typename batch_rule_t, batch_rule_t batch_rule>
  9216. at::Tensor relu_generated_plumbing(const at::Tensor & self) {
  9217. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9218. auto maybe_layer = maybeCurrentDynamicLayer();
  9219. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9220. int64_t cur_level = maybe_layer->layerId();
  9221. if (!isBatchedAtLevel(self, cur_level)) {
  9222. return at::_ops::relu::call(self);
  9223. }
  9224. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9225. auto results = batch_rule(self_value, self_bdim);
  9226. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9227. }
  9228. template <typename batch_rule_t, batch_rule_t batch_rule>
  9229. at::Tensor & relu__generated_plumbing(at::Tensor & self) {
  9230. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9231. auto maybe_layer = maybeCurrentDynamicLayer();
  9232. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  9233. int64_t cur_level = maybe_layer->layerId();
  9234. if (!isBatchedAtLevel(self, cur_level)) {
  9235. return at::_ops::relu_::call(self);
  9236. }
  9237. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9238. batch_rule(self_value, self_bdim);
  9239. return self;
  9240. }
  9241. template <typename batch_rule_t, batch_rule_t batch_rule>
  9242. at::Tensor relu6_generated_plumbing(const at::Tensor & self) {
  9243. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9244. auto maybe_layer = maybeCurrentDynamicLayer();
  9245. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9246. int64_t cur_level = maybe_layer->layerId();
  9247. if (!isBatchedAtLevel(self, cur_level)) {
  9248. return at::_ops::relu6::call(self);
  9249. }
  9250. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9251. auto results = batch_rule(self_value, self_bdim);
  9252. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9253. }
  9254. template <typename batch_rule_t, batch_rule_t batch_rule>
  9255. at::Tensor & relu6__generated_plumbing(at::Tensor & self) {
  9256. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9257. auto maybe_layer = maybeCurrentDynamicLayer();
  9258. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  9259. int64_t cur_level = maybe_layer->layerId();
  9260. if (!isBatchedAtLevel(self, cur_level)) {
  9261. return at::_ops::relu6_::call(self);
  9262. }
  9263. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9264. batch_rule(self_value, self_bdim);
  9265. return self;
  9266. }
  9267. template <typename batch_rule_t, batch_rule_t batch_rule>
  9268. at::Tensor prelu_generated_plumbing(const at::Tensor & self, const at::Tensor & weight) {
  9269. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9270. auto maybe_layer = maybeCurrentDynamicLayer();
  9271. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9272. int64_t cur_level = maybe_layer->layerId();
  9273. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(weight, cur_level)) {
  9274. return at::_ops::prelu::call(self, weight);
  9275. }
  9276. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9277. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  9278. auto results = batch_rule(self_value, self_bdim, weight_value, weight_bdim);
  9279. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9280. }
  9281. template <typename batch_rule_t, batch_rule_t batch_rule>
  9282. at::Tensor _prelu_kernel_generated_plumbing(const at::Tensor & self, const at::Tensor & weight) {
  9283. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9284. auto maybe_layer = maybeCurrentDynamicLayer();
  9285. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9286. int64_t cur_level = maybe_layer->layerId();
  9287. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(weight, cur_level)) {
  9288. return at::_ops::_prelu_kernel::call(self, weight);
  9289. }
  9290. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9291. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  9292. auto results = batch_rule(self_value, self_bdim, weight_value, weight_bdim);
  9293. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9294. }
  9295. template <typename batch_rule_t, batch_rule_t batch_rule>
  9296. ::std::tuple<at::Tensor,at::Tensor> _prelu_kernel_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self, const at::Tensor & weight) {
  9297. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9298. auto maybe_layer = maybeCurrentDynamicLayer();
  9299. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9300. int64_t cur_level = maybe_layer->layerId();
  9301. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(weight, cur_level)) {
  9302. return at::_ops::_prelu_kernel_backward::call(grad_output, self, weight);
  9303. }
  9304. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  9305. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9306. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  9307. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim, weight_value, weight_bdim);
  9308. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  9309. }
  9310. template <typename batch_rule_t, batch_rule_t batch_rule>
  9311. at::Tensor & gelu__generated_plumbing(at::Tensor & self, c10::string_view approximate) {
  9312. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9313. auto maybe_layer = maybeCurrentDynamicLayer();
  9314. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  9315. int64_t cur_level = maybe_layer->layerId();
  9316. if (!isBatchedAtLevel(self, cur_level)) {
  9317. return at::_ops::gelu_::call(self, approximate);
  9318. }
  9319. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9320. batch_rule(self_value, self_bdim, approximate);
  9321. return self;
  9322. }
  9323. template <typename batch_rule_t, batch_rule_t batch_rule>
  9324. at::Tensor gelu_generated_plumbing(const at::Tensor & self, c10::string_view approximate) {
  9325. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9326. auto maybe_layer = maybeCurrentDynamicLayer();
  9327. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9328. int64_t cur_level = maybe_layer->layerId();
  9329. if (!isBatchedAtLevel(self, cur_level)) {
  9330. return at::_ops::gelu::call(self, approximate);
  9331. }
  9332. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9333. auto results = batch_rule(self_value, self_bdim, approximate);
  9334. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9335. }
  9336. template <typename batch_rule_t, batch_rule_t batch_rule>
  9337. at::Tensor gelu_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self, c10::string_view approximate) {
  9338. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9339. auto maybe_layer = maybeCurrentDynamicLayer();
  9340. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9341. int64_t cur_level = maybe_layer->layerId();
  9342. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level)) {
  9343. return at::_ops::gelu_backward::call(grad_output, self, approximate);
  9344. }
  9345. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  9346. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9347. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim, approximate);
  9348. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9349. }
  9350. template <typename batch_rule_t, batch_rule_t batch_rule>
  9351. at::Tensor infinitely_differentiable_gelu_backward_generated_plumbing(const at::Tensor & grad, const at::Tensor & self) {
  9352. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9353. auto maybe_layer = maybeCurrentDynamicLayer();
  9354. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9355. int64_t cur_level = maybe_layer->layerId();
  9356. if (!isBatchedAtLevel(grad, cur_level) && !isBatchedAtLevel(self, cur_level)) {
  9357. return at::_ops::infinitely_differentiable_gelu_backward::call(grad, self);
  9358. }
  9359. auto [grad_value, grad_bdim] = unwrapTensorAtLevel(grad, cur_level);
  9360. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9361. auto results = batch_rule(grad_value, grad_bdim, self_value, self_bdim);
  9362. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9363. }
  9364. template <typename batch_rule_t, batch_rule_t batch_rule>
  9365. at::Tensor hardshrink_generated_plumbing(const at::Tensor & self, const at::Scalar & lambd) {
  9366. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9367. auto maybe_layer = maybeCurrentDynamicLayer();
  9368. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9369. int64_t cur_level = maybe_layer->layerId();
  9370. if (!isBatchedAtLevel(self, cur_level)) {
  9371. return at::_ops::hardshrink::call(self, lambd);
  9372. }
  9373. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9374. auto results = batch_rule(self_value, self_bdim, lambd);
  9375. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9376. }
  9377. template <typename batch_rule_t, batch_rule_t batch_rule>
  9378. at::Tensor hardshrink_backward_generated_plumbing(const at::Tensor & grad_out, const at::Tensor & self, const at::Scalar & lambd) {
  9379. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9380. auto maybe_layer = maybeCurrentDynamicLayer();
  9381. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9382. int64_t cur_level = maybe_layer->layerId();
  9383. if (!isBatchedAtLevel(grad_out, cur_level) && !isBatchedAtLevel(self, cur_level)) {
  9384. return at::_ops::hardshrink_backward::call(grad_out, self, lambd);
  9385. }
  9386. auto [grad_out_value, grad_out_bdim] = unwrapTensorAtLevel(grad_out, cur_level);
  9387. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9388. auto results = batch_rule(grad_out_value, grad_out_bdim, self_value, self_bdim, lambd);
  9389. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9390. }
  9391. template <typename batch_rule_t, batch_rule_t batch_rule>
  9392. at::Tensor rsqrt_generated_plumbing(const at::Tensor & self) {
  9393. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9394. auto maybe_layer = maybeCurrentDynamicLayer();
  9395. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9396. int64_t cur_level = maybe_layer->layerId();
  9397. if (!isBatchedAtLevel(self, cur_level)) {
  9398. return at::_ops::rsqrt::call(self);
  9399. }
  9400. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9401. auto results = batch_rule(self_value, self_bdim);
  9402. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9403. }
  9404. template <typename batch_rule_t, batch_rule_t batch_rule>
  9405. at::Tensor & rsqrt__generated_plumbing(at::Tensor & self) {
  9406. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9407. auto maybe_layer = maybeCurrentDynamicLayer();
  9408. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  9409. int64_t cur_level = maybe_layer->layerId();
  9410. if (!isBatchedAtLevel(self, cur_level)) {
  9411. return at::_ops::rsqrt_::call(self);
  9412. }
  9413. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9414. batch_rule(self_value, self_bdim);
  9415. return self;
  9416. }
  9417. template <typename batch_rule_t, batch_rule_t batch_rule>
  9418. at::Tensor select_Dimname_generated_plumbing(const at::Tensor & self, at::Dimname dim, int64_t index) {
  9419. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9420. auto maybe_layer = maybeCurrentDynamicLayer();
  9421. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9422. int64_t cur_level = maybe_layer->layerId();
  9423. if (!isBatchedAtLevel(self, cur_level)) {
  9424. return at::_ops::select_Dimname::call(self, dim, index);
  9425. }
  9426. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9427. auto results = batch_rule(self_value, self_bdim, dim, index);
  9428. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9429. }
  9430. template <typename batch_rule_t, batch_rule_t batch_rule>
  9431. at::Tensor select_int_generated_plumbing(const at::Tensor & self, int64_t dim, c10::SymInt index) {
  9432. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9433. auto maybe_layer = maybeCurrentDynamicLayer();
  9434. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9435. int64_t cur_level = maybe_layer->layerId();
  9436. if (!isBatchedAtLevel(self, cur_level)) {
  9437. return at::_ops::select_int::call(self, dim, index);
  9438. }
  9439. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9440. auto results = batch_rule(self_value, self_bdim, dim, index);
  9441. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9442. }
  9443. template <typename batch_rule_t, batch_rule_t batch_rule>
  9444. at::Tensor select_backward_generated_plumbing(const at::Tensor & grad_output, c10::SymIntArrayRef input_sizes, int64_t dim, c10::SymInt index) {
  9445. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9446. auto maybe_layer = maybeCurrentDynamicLayer();
  9447. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9448. int64_t cur_level = maybe_layer->layerId();
  9449. if (!isBatchedAtLevel(grad_output, cur_level)) {
  9450. return at::_ops::select_backward::call(grad_output, input_sizes, dim, index);
  9451. }
  9452. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  9453. auto results = batch_rule(grad_output_value, grad_output_bdim, input_sizes, dim, index);
  9454. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9455. }
  9456. template <typename batch_rule_t, batch_rule_t batch_rule>
  9457. at::Tensor _nested_select_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self, int64_t dim, c10::SymInt index) {
  9458. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9459. auto maybe_layer = maybeCurrentDynamicLayer();
  9460. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9461. int64_t cur_level = maybe_layer->layerId();
  9462. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level)) {
  9463. return at::_ops::_nested_select_backward::call(grad_output, self, dim, index);
  9464. }
  9465. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  9466. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9467. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim, dim, index);
  9468. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9469. }
  9470. template <typename batch_rule_t, batch_rule_t batch_rule>
  9471. at::Tensor selu_generated_plumbing(const at::Tensor & self) {
  9472. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9473. auto maybe_layer = maybeCurrentDynamicLayer();
  9474. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9475. int64_t cur_level = maybe_layer->layerId();
  9476. if (!isBatchedAtLevel(self, cur_level)) {
  9477. return at::_ops::selu::call(self);
  9478. }
  9479. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9480. auto results = batch_rule(self_value, self_bdim);
  9481. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9482. }
  9483. template <typename batch_rule_t, batch_rule_t batch_rule>
  9484. at::Tensor & selu__generated_plumbing(at::Tensor & self) {
  9485. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9486. auto maybe_layer = maybeCurrentDynamicLayer();
  9487. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  9488. int64_t cur_level = maybe_layer->layerId();
  9489. if (!isBatchedAtLevel(self, cur_level)) {
  9490. return at::_ops::selu_::call(self);
  9491. }
  9492. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9493. batch_rule(self_value, self_bdim);
  9494. return self;
  9495. }
  9496. template <typename batch_rule_t, batch_rule_t batch_rule>
  9497. at::Tensor celu_generated_plumbing(const at::Tensor & self, const at::Scalar & alpha) {
  9498. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9499. auto maybe_layer = maybeCurrentDynamicLayer();
  9500. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9501. int64_t cur_level = maybe_layer->layerId();
  9502. if (!isBatchedAtLevel(self, cur_level)) {
  9503. return at::_ops::celu::call(self, alpha);
  9504. }
  9505. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9506. auto results = batch_rule(self_value, self_bdim, alpha);
  9507. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9508. }
  9509. template <typename batch_rule_t, batch_rule_t batch_rule>
  9510. at::Tensor & celu__generated_plumbing(at::Tensor & self, const at::Scalar & alpha) {
  9511. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9512. auto maybe_layer = maybeCurrentDynamicLayer();
  9513. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  9514. int64_t cur_level = maybe_layer->layerId();
  9515. if (!isBatchedAtLevel(self, cur_level)) {
  9516. return at::_ops::celu_::call(self, alpha);
  9517. }
  9518. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9519. batch_rule(self_value, self_bdim, alpha);
  9520. return self;
  9521. }
  9522. template <typename batch_rule_t, batch_rule_t batch_rule>
  9523. at::Tensor silu_generated_plumbing(const at::Tensor & self) {
  9524. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9525. auto maybe_layer = maybeCurrentDynamicLayer();
  9526. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9527. int64_t cur_level = maybe_layer->layerId();
  9528. if (!isBatchedAtLevel(self, cur_level)) {
  9529. return at::_ops::silu::call(self);
  9530. }
  9531. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9532. auto results = batch_rule(self_value, self_bdim);
  9533. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9534. }
  9535. template <typename batch_rule_t, batch_rule_t batch_rule>
  9536. at::Tensor & silu__generated_plumbing(at::Tensor & self) {
  9537. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9538. auto maybe_layer = maybeCurrentDynamicLayer();
  9539. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  9540. int64_t cur_level = maybe_layer->layerId();
  9541. if (!isBatchedAtLevel(self, cur_level)) {
  9542. return at::_ops::silu_::call(self);
  9543. }
  9544. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9545. batch_rule(self_value, self_bdim);
  9546. return self;
  9547. }
  9548. template <typename batch_rule_t, batch_rule_t batch_rule>
  9549. at::Tensor silu_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self) {
  9550. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9551. auto maybe_layer = maybeCurrentDynamicLayer();
  9552. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9553. int64_t cur_level = maybe_layer->layerId();
  9554. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level)) {
  9555. return at::_ops::silu_backward::call(grad_output, self);
  9556. }
  9557. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  9558. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9559. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim);
  9560. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9561. }
  9562. template <typename batch_rule_t, batch_rule_t batch_rule>
  9563. at::Tensor mish_generated_plumbing(const at::Tensor & self) {
  9564. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9565. auto maybe_layer = maybeCurrentDynamicLayer();
  9566. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9567. int64_t cur_level = maybe_layer->layerId();
  9568. if (!isBatchedAtLevel(self, cur_level)) {
  9569. return at::_ops::mish::call(self);
  9570. }
  9571. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9572. auto results = batch_rule(self_value, self_bdim);
  9573. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9574. }
  9575. template <typename batch_rule_t, batch_rule_t batch_rule>
  9576. at::Tensor & mish__generated_plumbing(at::Tensor & self) {
  9577. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9578. auto maybe_layer = maybeCurrentDynamicLayer();
  9579. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  9580. int64_t cur_level = maybe_layer->layerId();
  9581. if (!isBatchedAtLevel(self, cur_level)) {
  9582. return at::_ops::mish_::call(self);
  9583. }
  9584. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9585. batch_rule(self_value, self_bdim);
  9586. return self;
  9587. }
  9588. template <typename batch_rule_t, batch_rule_t batch_rule>
  9589. at::Tensor mish_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self) {
  9590. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9591. auto maybe_layer = maybeCurrentDynamicLayer();
  9592. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9593. int64_t cur_level = maybe_layer->layerId();
  9594. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level)) {
  9595. return at::_ops::mish_backward::call(grad_output, self);
  9596. }
  9597. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  9598. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9599. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim);
  9600. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9601. }
  9602. template <typename batch_rule_t, batch_rule_t batch_rule>
  9603. at::Tensor sigmoid_generated_plumbing(const at::Tensor & self) {
  9604. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9605. auto maybe_layer = maybeCurrentDynamicLayer();
  9606. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9607. int64_t cur_level = maybe_layer->layerId();
  9608. if (!isBatchedAtLevel(self, cur_level)) {
  9609. return at::_ops::sigmoid::call(self);
  9610. }
  9611. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9612. auto results = batch_rule(self_value, self_bdim);
  9613. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9614. }
  9615. template <typename batch_rule_t, batch_rule_t batch_rule>
  9616. at::Tensor & sigmoid__generated_plumbing(at::Tensor & self) {
  9617. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9618. auto maybe_layer = maybeCurrentDynamicLayer();
  9619. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  9620. int64_t cur_level = maybe_layer->layerId();
  9621. if (!isBatchedAtLevel(self, cur_level)) {
  9622. return at::_ops::sigmoid_::call(self);
  9623. }
  9624. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9625. batch_rule(self_value, self_bdim);
  9626. return self;
  9627. }
  9628. template <typename batch_rule_t, batch_rule_t batch_rule>
  9629. at::Tensor logit_generated_plumbing(const at::Tensor & self, ::std::optional<double> eps) {
  9630. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9631. auto maybe_layer = maybeCurrentDynamicLayer();
  9632. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9633. int64_t cur_level = maybe_layer->layerId();
  9634. if (!isBatchedAtLevel(self, cur_level)) {
  9635. return at::_ops::logit::call(self, eps);
  9636. }
  9637. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9638. auto results = batch_rule(self_value, self_bdim, eps);
  9639. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9640. }
  9641. template <typename batch_rule_t, batch_rule_t batch_rule>
  9642. at::Tensor & logit__generated_plumbing(at::Tensor & self, ::std::optional<double> eps) {
  9643. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9644. auto maybe_layer = maybeCurrentDynamicLayer();
  9645. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  9646. int64_t cur_level = maybe_layer->layerId();
  9647. if (!isBatchedAtLevel(self, cur_level)) {
  9648. return at::_ops::logit_::call(self, eps);
  9649. }
  9650. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9651. batch_rule(self_value, self_bdim, eps);
  9652. return self;
  9653. }
  9654. template <typename batch_rule_t, batch_rule_t batch_rule>
  9655. at::Tensor sin_generated_plumbing(const at::Tensor & self) {
  9656. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9657. auto maybe_layer = maybeCurrentDynamicLayer();
  9658. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9659. int64_t cur_level = maybe_layer->layerId();
  9660. if (!isBatchedAtLevel(self, cur_level)) {
  9661. return at::_ops::sin::call(self);
  9662. }
  9663. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9664. auto results = batch_rule(self_value, self_bdim);
  9665. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9666. }
  9667. template <typename batch_rule_t, batch_rule_t batch_rule>
  9668. at::Tensor & sin__generated_plumbing(at::Tensor & self) {
  9669. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9670. auto maybe_layer = maybeCurrentDynamicLayer();
  9671. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  9672. int64_t cur_level = maybe_layer->layerId();
  9673. if (!isBatchedAtLevel(self, cur_level)) {
  9674. return at::_ops::sin_::call(self);
  9675. }
  9676. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9677. batch_rule(self_value, self_bdim);
  9678. return self;
  9679. }
  9680. template <typename batch_rule_t, batch_rule_t batch_rule>
  9681. at::Tensor sinc_generated_plumbing(const at::Tensor & self) {
  9682. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9683. auto maybe_layer = maybeCurrentDynamicLayer();
  9684. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9685. int64_t cur_level = maybe_layer->layerId();
  9686. if (!isBatchedAtLevel(self, cur_level)) {
  9687. return at::_ops::sinc::call(self);
  9688. }
  9689. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9690. auto results = batch_rule(self_value, self_bdim);
  9691. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9692. }
  9693. template <typename batch_rule_t, batch_rule_t batch_rule>
  9694. at::Tensor & sinc__generated_plumbing(at::Tensor & self) {
  9695. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9696. auto maybe_layer = maybeCurrentDynamicLayer();
  9697. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  9698. int64_t cur_level = maybe_layer->layerId();
  9699. if (!isBatchedAtLevel(self, cur_level)) {
  9700. return at::_ops::sinc_::call(self);
  9701. }
  9702. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9703. batch_rule(self_value, self_bdim);
  9704. return self;
  9705. }
  9706. template <typename batch_rule_t, batch_rule_t batch_rule>
  9707. at::Tensor sinh_generated_plumbing(const at::Tensor & self) {
  9708. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9709. auto maybe_layer = maybeCurrentDynamicLayer();
  9710. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9711. int64_t cur_level = maybe_layer->layerId();
  9712. if (!isBatchedAtLevel(self, cur_level)) {
  9713. return at::_ops::sinh::call(self);
  9714. }
  9715. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9716. auto results = batch_rule(self_value, self_bdim);
  9717. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9718. }
  9719. template <typename batch_rule_t, batch_rule_t batch_rule>
  9720. at::Tensor & sinh__generated_plumbing(at::Tensor & self) {
  9721. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9722. auto maybe_layer = maybeCurrentDynamicLayer();
  9723. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  9724. int64_t cur_level = maybe_layer->layerId();
  9725. if (!isBatchedAtLevel(self, cur_level)) {
  9726. return at::_ops::sinh_::call(self);
  9727. }
  9728. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9729. batch_rule(self_value, self_bdim);
  9730. return self;
  9731. }
  9732. template <typename batch_rule_t, batch_rule_t batch_rule>
  9733. at::Tensor detach_generated_plumbing(const at::Tensor & self) {
  9734. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9735. auto maybe_layer = maybeCurrentDynamicLayer();
  9736. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9737. int64_t cur_level = maybe_layer->layerId();
  9738. if (!isBatchedAtLevel(self, cur_level)) {
  9739. return at::_ops::detach::call(self);
  9740. }
  9741. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9742. auto results = batch_rule(self_value, self_bdim);
  9743. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9744. }
  9745. template <typename batch_rule_t, batch_rule_t batch_rule>
  9746. at::Tensor slice_Tensor_generated_plumbing(const at::Tensor & self, int64_t dim, ::std::optional<c10::SymInt> start, ::std::optional<c10::SymInt> end, c10::SymInt step) {
  9747. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9748. auto maybe_layer = maybeCurrentDynamicLayer();
  9749. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9750. int64_t cur_level = maybe_layer->layerId();
  9751. if (!isBatchedAtLevel(self, cur_level)) {
  9752. return at::_ops::slice_Tensor::call(self, dim, start, end, step);
  9753. }
  9754. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9755. auto results = batch_rule(self_value, self_bdim, dim, start, end, step);
  9756. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9757. }
  9758. template <typename batch_rule_t, batch_rule_t batch_rule>
  9759. at::Tensor slice_backward_generated_plumbing(const at::Tensor & grad_output, c10::SymIntArrayRef input_sizes, int64_t dim, c10::SymInt start, c10::SymInt end, c10::SymInt step) {
  9760. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9761. auto maybe_layer = maybeCurrentDynamicLayer();
  9762. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9763. int64_t cur_level = maybe_layer->layerId();
  9764. if (!isBatchedAtLevel(grad_output, cur_level)) {
  9765. return at::_ops::slice_backward::call(grad_output, input_sizes, dim, start, end, step);
  9766. }
  9767. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  9768. auto results = batch_rule(grad_output_value, grad_output_bdim, input_sizes, dim, start, end, step);
  9769. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9770. }
  9771. template <typename batch_rule_t, batch_rule_t batch_rule>
  9772. at::Tensor slice_inverse_generated_plumbing(const at::Tensor & self, const at::Tensor & src, int64_t dim, ::std::optional<c10::SymInt> start, ::std::optional<c10::SymInt> end, c10::SymInt step) {
  9773. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9774. auto maybe_layer = maybeCurrentDynamicLayer();
  9775. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9776. int64_t cur_level = maybe_layer->layerId();
  9777. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(src, cur_level)) {
  9778. return at::_ops::slice_inverse::call(self, src, dim, start, end, step);
  9779. }
  9780. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9781. auto [src_value, src_bdim] = unwrapTensorAtLevel(src, cur_level);
  9782. auto results = batch_rule(self_value, self_bdim, src_value, src_bdim, dim, start, end, step);
  9783. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9784. }
  9785. template <typename batch_rule_t, batch_rule_t batch_rule>
  9786. at::Tensor slice_scatter_generated_plumbing(const at::Tensor & self, const at::Tensor & src, int64_t dim, ::std::optional<c10::SymInt> start, ::std::optional<c10::SymInt> end, c10::SymInt step) {
  9787. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9788. auto maybe_layer = maybeCurrentDynamicLayer();
  9789. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9790. int64_t cur_level = maybe_layer->layerId();
  9791. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(src, cur_level)) {
  9792. return at::_ops::slice_scatter::call(self, src, dim, start, end, step);
  9793. }
  9794. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9795. auto [src_value, src_bdim] = unwrapTensorAtLevel(src, cur_level);
  9796. auto results = batch_rule(self_value, self_bdim, src_value, src_bdim, dim, start, end, step);
  9797. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9798. }
  9799. template <typename batch_rule_t, batch_rule_t batch_rule>
  9800. at::Tensor select_scatter_generated_plumbing(const at::Tensor & self, const at::Tensor & src, int64_t dim, c10::SymInt index) {
  9801. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9802. auto maybe_layer = maybeCurrentDynamicLayer();
  9803. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9804. int64_t cur_level = maybe_layer->layerId();
  9805. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(src, cur_level)) {
  9806. return at::_ops::select_scatter::call(self, src, dim, index);
  9807. }
  9808. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9809. auto [src_value, src_bdim] = unwrapTensorAtLevel(src, cur_level);
  9810. auto results = batch_rule(self_value, self_bdim, src_value, src_bdim, dim, index);
  9811. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9812. }
  9813. template <typename batch_rule_t, batch_rule_t batch_rule>
  9814. at::Tensor diagonal_scatter_generated_plumbing(const at::Tensor & self, const at::Tensor & src, int64_t offset, int64_t dim1, int64_t dim2) {
  9815. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9816. auto maybe_layer = maybeCurrentDynamicLayer();
  9817. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9818. int64_t cur_level = maybe_layer->layerId();
  9819. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(src, cur_level)) {
  9820. return at::_ops::diagonal_scatter::call(self, src, offset, dim1, dim2);
  9821. }
  9822. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9823. auto [src_value, src_bdim] = unwrapTensorAtLevel(src, cur_level);
  9824. auto results = batch_rule(self_value, self_bdim, src_value, src_bdim, offset, dim1, dim2);
  9825. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9826. }
  9827. template <typename batch_rule_t, batch_rule_t batch_rule>
  9828. at::Tensor as_strided_scatter_generated_plumbing(const at::Tensor & self, const at::Tensor & src, c10::SymIntArrayRef size, c10::SymIntArrayRef stride, ::std::optional<c10::SymInt> storage_offset) {
  9829. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9830. auto maybe_layer = maybeCurrentDynamicLayer();
  9831. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9832. int64_t cur_level = maybe_layer->layerId();
  9833. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(src, cur_level)) {
  9834. return at::_ops::as_strided_scatter::call(self, src, size, stride, storage_offset);
  9835. }
  9836. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9837. auto [src_value, src_bdim] = unwrapTensorAtLevel(src, cur_level);
  9838. auto results = batch_rule(self_value, self_bdim, src_value, src_bdim, size, stride, storage_offset);
  9839. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9840. }
  9841. template <typename batch_rule_t, batch_rule_t batch_rule>
  9842. at::Tensor smm_generated_plumbing(const at::Tensor & self, const at::Tensor & mat2) {
  9843. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9844. auto maybe_layer = maybeCurrentDynamicLayer();
  9845. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9846. int64_t cur_level = maybe_layer->layerId();
  9847. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(mat2, cur_level)) {
  9848. return at::_ops::smm::call(self, mat2);
  9849. }
  9850. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9851. auto [mat2_value, mat2_bdim] = unwrapTensorAtLevel(mat2, cur_level);
  9852. auto results = batch_rule(self_value, self_bdim, mat2_value, mat2_bdim);
  9853. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9854. }
  9855. template <typename batch_rule_t, batch_rule_t batch_rule>
  9856. at::Tensor softmax_int_generated_plumbing(const at::Tensor & self, int64_t dim, ::std::optional<at::ScalarType> dtype) {
  9857. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9858. auto maybe_layer = maybeCurrentDynamicLayer();
  9859. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9860. int64_t cur_level = maybe_layer->layerId();
  9861. if (!isBatchedAtLevel(self, cur_level)) {
  9862. return at::_ops::softmax_int::call(self, dim, dtype);
  9863. }
  9864. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9865. auto results = batch_rule(self_value, self_bdim, dim, dtype);
  9866. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9867. }
  9868. template <typename batch_rule_t, batch_rule_t batch_rule>
  9869. at::Tensor softmax_Dimname_generated_plumbing(const at::Tensor & self, at::Dimname dim, ::std::optional<at::ScalarType> dtype) {
  9870. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9871. auto maybe_layer = maybeCurrentDynamicLayer();
  9872. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9873. int64_t cur_level = maybe_layer->layerId();
  9874. if (!isBatchedAtLevel(self, cur_level)) {
  9875. return at::_ops::softmax_Dimname::call(self, dim, dtype);
  9876. }
  9877. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9878. auto results = batch_rule(self_value, self_bdim, dim, dtype);
  9879. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9880. }
  9881. template <typename batch_rule_t, batch_rule_t batch_rule>
  9882. at::Tensor _softmax_generated_plumbing(const at::Tensor & self, int64_t dim, bool half_to_float) {
  9883. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9884. auto maybe_layer = maybeCurrentDynamicLayer();
  9885. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9886. int64_t cur_level = maybe_layer->layerId();
  9887. if (!isBatchedAtLevel(self, cur_level)) {
  9888. return at::_ops::_softmax::call(self, dim, half_to_float);
  9889. }
  9890. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9891. auto results = batch_rule(self_value, self_bdim, dim, half_to_float);
  9892. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9893. }
  9894. template <typename batch_rule_t, batch_rule_t batch_rule>
  9895. at::Tensor _softmax_backward_data_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & output, int64_t dim, at::ScalarType input_dtype) {
  9896. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9897. auto maybe_layer = maybeCurrentDynamicLayer();
  9898. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9899. int64_t cur_level = maybe_layer->layerId();
  9900. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(output, cur_level)) {
  9901. return at::_ops::_softmax_backward_data::call(grad_output, output, dim, input_dtype);
  9902. }
  9903. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  9904. auto [output_value, output_bdim] = unwrapTensorAtLevel(output, cur_level);
  9905. auto results = batch_rule(grad_output_value, grad_output_bdim, output_value, output_bdim, dim, input_dtype);
  9906. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  9907. }
  9908. template <typename batch_rule_t, batch_rule_t batch_rule>
  9909. ::std::vector<at::Tensor> unsafe_split_Tensor_generated_plumbing(const at::Tensor & self, c10::SymInt split_size, int64_t dim) {
  9910. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9911. auto maybe_layer = maybeCurrentDynamicLayer();
  9912. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9913. int64_t cur_level = maybe_layer->layerId();
  9914. if (!isBatchedAtLevel(self, cur_level)) {
  9915. return at::_ops::unsafe_split_Tensor::call(self, split_size, dim);
  9916. }
  9917. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9918. auto results = batch_rule(self_value, self_bdim, split_size, dim);
  9919. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  9920. }
  9921. template <typename batch_rule_t, batch_rule_t batch_rule>
  9922. ::std::vector<at::Tensor> split_Tensor_generated_plumbing(const at::Tensor & self, c10::SymInt split_size, int64_t dim) {
  9923. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9924. auto maybe_layer = maybeCurrentDynamicLayer();
  9925. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9926. int64_t cur_level = maybe_layer->layerId();
  9927. if (!isBatchedAtLevel(self, cur_level)) {
  9928. return at::_ops::split_Tensor::call(self, split_size, dim);
  9929. }
  9930. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9931. auto results = batch_rule(self_value, self_bdim, split_size, dim);
  9932. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  9933. }
  9934. template <typename batch_rule_t, batch_rule_t batch_rule>
  9935. ::std::vector<at::Tensor> split_sizes_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef split_size, int64_t dim) {
  9936. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9937. auto maybe_layer = maybeCurrentDynamicLayer();
  9938. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9939. int64_t cur_level = maybe_layer->layerId();
  9940. if (!isBatchedAtLevel(self, cur_level)) {
  9941. return at::_ops::split_sizes::call(self, split_size, dim);
  9942. }
  9943. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9944. auto results = batch_rule(self_value, self_bdim, split_size, dim);
  9945. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  9946. }
  9947. template <typename batch_rule_t, batch_rule_t batch_rule>
  9948. ::std::vector<at::Tensor> unsafe_split_with_sizes_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef split_sizes, int64_t dim) {
  9949. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9950. auto maybe_layer = maybeCurrentDynamicLayer();
  9951. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9952. int64_t cur_level = maybe_layer->layerId();
  9953. if (!isBatchedAtLevel(self, cur_level)) {
  9954. return at::_ops::unsafe_split_with_sizes::call(self, split_sizes, dim);
  9955. }
  9956. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9957. auto results = batch_rule(self_value, self_bdim, split_sizes, dim);
  9958. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  9959. }
  9960. template <typename batch_rule_t, batch_rule_t batch_rule>
  9961. ::std::vector<at::Tensor> split_with_sizes_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef split_sizes, int64_t dim) {
  9962. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9963. auto maybe_layer = maybeCurrentDynamicLayer();
  9964. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9965. int64_t cur_level = maybe_layer->layerId();
  9966. if (!isBatchedAtLevel(self, cur_level)) {
  9967. return at::_ops::split_with_sizes::call(self, split_sizes, dim);
  9968. }
  9969. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9970. auto results = batch_rule(self_value, self_bdim, split_sizes, dim);
  9971. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  9972. }
  9973. template <typename batch_rule_t, batch_rule_t batch_rule>
  9974. ::std::vector<at::Tensor> hsplit_int_generated_plumbing(const at::Tensor & self, int64_t sections) {
  9975. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9976. auto maybe_layer = maybeCurrentDynamicLayer();
  9977. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9978. int64_t cur_level = maybe_layer->layerId();
  9979. if (!isBatchedAtLevel(self, cur_level)) {
  9980. return at::_ops::hsplit_int::call(self, sections);
  9981. }
  9982. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9983. auto results = batch_rule(self_value, self_bdim, sections);
  9984. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  9985. }
  9986. template <typename batch_rule_t, batch_rule_t batch_rule>
  9987. ::std::vector<at::Tensor> hsplit_array_generated_plumbing(const at::Tensor & self, at::IntArrayRef indices) {
  9988. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  9989. auto maybe_layer = maybeCurrentDynamicLayer();
  9990. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  9991. int64_t cur_level = maybe_layer->layerId();
  9992. if (!isBatchedAtLevel(self, cur_level)) {
  9993. return at::_ops::hsplit_array::call(self, indices);
  9994. }
  9995. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  9996. auto results = batch_rule(self_value, self_bdim, indices);
  9997. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  9998. }
  9999. template <typename batch_rule_t, batch_rule_t batch_rule>
  10000. ::std::vector<at::Tensor> vsplit_int_generated_plumbing(const at::Tensor & self, int64_t sections) {
  10001. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10002. auto maybe_layer = maybeCurrentDynamicLayer();
  10003. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10004. int64_t cur_level = maybe_layer->layerId();
  10005. if (!isBatchedAtLevel(self, cur_level)) {
  10006. return at::_ops::vsplit_int::call(self, sections);
  10007. }
  10008. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10009. auto results = batch_rule(self_value, self_bdim, sections);
  10010. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  10011. }
  10012. template <typename batch_rule_t, batch_rule_t batch_rule>
  10013. ::std::vector<at::Tensor> vsplit_array_generated_plumbing(const at::Tensor & self, at::IntArrayRef indices) {
  10014. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10015. auto maybe_layer = maybeCurrentDynamicLayer();
  10016. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10017. int64_t cur_level = maybe_layer->layerId();
  10018. if (!isBatchedAtLevel(self, cur_level)) {
  10019. return at::_ops::vsplit_array::call(self, indices);
  10020. }
  10021. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10022. auto results = batch_rule(self_value, self_bdim, indices);
  10023. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  10024. }
  10025. template <typename batch_rule_t, batch_rule_t batch_rule>
  10026. ::std::vector<at::Tensor> dsplit_int_generated_plumbing(const at::Tensor & self, int64_t sections) {
  10027. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10028. auto maybe_layer = maybeCurrentDynamicLayer();
  10029. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10030. int64_t cur_level = maybe_layer->layerId();
  10031. if (!isBatchedAtLevel(self, cur_level)) {
  10032. return at::_ops::dsplit_int::call(self, sections);
  10033. }
  10034. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10035. auto results = batch_rule(self_value, self_bdim, sections);
  10036. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  10037. }
  10038. template <typename batch_rule_t, batch_rule_t batch_rule>
  10039. ::std::vector<at::Tensor> dsplit_array_generated_plumbing(const at::Tensor & self, at::IntArrayRef indices) {
  10040. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10041. auto maybe_layer = maybeCurrentDynamicLayer();
  10042. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10043. int64_t cur_level = maybe_layer->layerId();
  10044. if (!isBatchedAtLevel(self, cur_level)) {
  10045. return at::_ops::dsplit_array::call(self, indices);
  10046. }
  10047. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10048. auto results = batch_rule(self_value, self_bdim, indices);
  10049. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  10050. }
  10051. template <typename batch_rule_t, batch_rule_t batch_rule>
  10052. at::Tensor squeeze_generated_plumbing(const at::Tensor & self) {
  10053. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10054. auto maybe_layer = maybeCurrentDynamicLayer();
  10055. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10056. int64_t cur_level = maybe_layer->layerId();
  10057. if (!isBatchedAtLevel(self, cur_level)) {
  10058. return at::_ops::squeeze::call(self);
  10059. }
  10060. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10061. auto results = batch_rule(self_value, self_bdim);
  10062. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10063. }
  10064. template <typename batch_rule_t, batch_rule_t batch_rule>
  10065. at::Tensor squeeze_dim_generated_plumbing(const at::Tensor & self, int64_t dim) {
  10066. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10067. auto maybe_layer = maybeCurrentDynamicLayer();
  10068. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10069. int64_t cur_level = maybe_layer->layerId();
  10070. if (!isBatchedAtLevel(self, cur_level)) {
  10071. return at::_ops::squeeze_dim::call(self, dim);
  10072. }
  10073. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10074. auto results = batch_rule(self_value, self_bdim, dim);
  10075. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10076. }
  10077. template <typename batch_rule_t, batch_rule_t batch_rule>
  10078. at::Tensor squeeze_dimname_generated_plumbing(const at::Tensor & self, at::Dimname dim) {
  10079. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10080. auto maybe_layer = maybeCurrentDynamicLayer();
  10081. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10082. int64_t cur_level = maybe_layer->layerId();
  10083. if (!isBatchedAtLevel(self, cur_level)) {
  10084. return at::_ops::squeeze_dimname::call(self, dim);
  10085. }
  10086. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10087. auto results = batch_rule(self_value, self_bdim, dim);
  10088. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10089. }
  10090. template <typename batch_rule_t, batch_rule_t batch_rule>
  10091. at::Tensor squeeze_dims_generated_plumbing(const at::Tensor & self, at::IntArrayRef dim) {
  10092. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10093. auto maybe_layer = maybeCurrentDynamicLayer();
  10094. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10095. int64_t cur_level = maybe_layer->layerId();
  10096. if (!isBatchedAtLevel(self, cur_level)) {
  10097. return at::_ops::squeeze_dims::call(self, dim);
  10098. }
  10099. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10100. auto results = batch_rule(self_value, self_bdim, dim);
  10101. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10102. }
  10103. template <typename batch_rule_t, batch_rule_t batch_rule>
  10104. at::Tensor sspaddmm_generated_plumbing(const at::Tensor & self, const at::Tensor & mat1, const at::Tensor & mat2, const at::Scalar & beta, const at::Scalar & alpha) {
  10105. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10106. auto maybe_layer = maybeCurrentDynamicLayer();
  10107. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10108. int64_t cur_level = maybe_layer->layerId();
  10109. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(mat1, cur_level) && !isBatchedAtLevel(mat2, cur_level)) {
  10110. return at::_ops::sspaddmm::call(self, mat1, mat2, beta, alpha);
  10111. }
  10112. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10113. auto [mat1_value, mat1_bdim] = unwrapTensorAtLevel(mat1, cur_level);
  10114. auto [mat2_value, mat2_bdim] = unwrapTensorAtLevel(mat2, cur_level);
  10115. auto results = batch_rule(self_value, self_bdim, mat1_value, mat1_bdim, mat2_value, mat2_bdim, beta, alpha);
  10116. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10117. }
  10118. template <typename batch_rule_t, batch_rule_t batch_rule>
  10119. at::Tensor _chunk_cat_generated_plumbing(at::TensorList tensors, int64_t dim, int64_t num_chunks) {
  10120. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10121. auto maybe_layer = maybeCurrentDynamicLayer();
  10122. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10123. int64_t cur_level = maybe_layer->layerId();
  10124. if (!isBatchedAtLevel(tensors, cur_level)) {
  10125. return at::_ops::_chunk_cat::call(tensors, dim, num_chunks);
  10126. }
  10127. auto results = batch_rule(tensors, dim, num_chunks);
  10128. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10129. }
  10130. template <typename batch_rule_t, batch_rule_t batch_rule>
  10131. at::Tensor stack_generated_plumbing(at::TensorList tensors, int64_t dim) {
  10132. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10133. auto maybe_layer = maybeCurrentDynamicLayer();
  10134. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10135. int64_t cur_level = maybe_layer->layerId();
  10136. if (!isBatchedAtLevel(tensors, cur_level)) {
  10137. return at::_ops::stack::call(tensors, dim);
  10138. }
  10139. auto results = batch_rule(tensors, dim);
  10140. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10141. }
  10142. template <typename batch_rule_t, batch_rule_t batch_rule>
  10143. at::Tensor _stack_generated_plumbing(at::TensorList tensors, int64_t dim) {
  10144. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10145. auto maybe_layer = maybeCurrentDynamicLayer();
  10146. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10147. int64_t cur_level = maybe_layer->layerId();
  10148. if (!isBatchedAtLevel(tensors, cur_level)) {
  10149. return at::_ops::_stack::call(tensors, dim);
  10150. }
  10151. auto results = batch_rule(tensors, dim);
  10152. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10153. }
  10154. template <typename batch_rule_t, batch_rule_t batch_rule>
  10155. at::Tensor hstack_generated_plumbing(at::TensorList tensors) {
  10156. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10157. auto maybe_layer = maybeCurrentDynamicLayer();
  10158. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10159. int64_t cur_level = maybe_layer->layerId();
  10160. if (!isBatchedAtLevel(tensors, cur_level)) {
  10161. return at::_ops::hstack::call(tensors);
  10162. }
  10163. auto results = batch_rule(tensors);
  10164. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10165. }
  10166. template <typename batch_rule_t, batch_rule_t batch_rule>
  10167. at::Tensor vstack_generated_plumbing(at::TensorList tensors) {
  10168. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10169. auto maybe_layer = maybeCurrentDynamicLayer();
  10170. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10171. int64_t cur_level = maybe_layer->layerId();
  10172. if (!isBatchedAtLevel(tensors, cur_level)) {
  10173. return at::_ops::vstack::call(tensors);
  10174. }
  10175. auto results = batch_rule(tensors);
  10176. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10177. }
  10178. template <typename batch_rule_t, batch_rule_t batch_rule>
  10179. at::Tensor dstack_generated_plumbing(at::TensorList tensors) {
  10180. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10181. auto maybe_layer = maybeCurrentDynamicLayer();
  10182. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10183. int64_t cur_level = maybe_layer->layerId();
  10184. if (!isBatchedAtLevel(tensors, cur_level)) {
  10185. return at::_ops::dstack::call(tensors);
  10186. }
  10187. auto results = batch_rule(tensors);
  10188. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10189. }
  10190. template <typename batch_rule_t, batch_rule_t batch_rule>
  10191. at::Tensor stft_generated_plumbing(const at::Tensor & self, int64_t n_fft, ::std::optional<int64_t> hop_length, ::std::optional<int64_t> win_length, const ::std::optional<at::Tensor> & window, bool normalized, ::std::optional<bool> onesided, ::std::optional<bool> return_complex, ::std::optional<bool> align_to_window) {
  10192. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10193. auto maybe_layer = maybeCurrentDynamicLayer();
  10194. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10195. int64_t cur_level = maybe_layer->layerId();
  10196. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(window, cur_level)) {
  10197. return at::_ops::stft::call(self, n_fft, hop_length, win_length, window, normalized, onesided, return_complex, align_to_window);
  10198. }
  10199. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10200. std::optional<Tensor> window_value;
  10201. std::optional<int64_t> window_bdim;
  10202. if (window) {
  10203. std::tie(window_value, window_bdim) = unwrapTensorAtLevel(window.value(), cur_level);
  10204. }
  10205. auto results = batch_rule(self_value, self_bdim, n_fft, hop_length, win_length, window_value, window_bdim, normalized, onesided, return_complex, align_to_window);
  10206. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10207. }
  10208. template <typename batch_rule_t, batch_rule_t batch_rule>
  10209. at::Tensor stft_center_generated_plumbing(const at::Tensor & self, int64_t n_fft, ::std::optional<int64_t> hop_length, ::std::optional<int64_t> win_length, const ::std::optional<at::Tensor> & window, bool center, c10::string_view pad_mode, bool normalized, ::std::optional<bool> onesided, ::std::optional<bool> return_complex, ::std::optional<bool> align_to_window) {
  10210. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10211. auto maybe_layer = maybeCurrentDynamicLayer();
  10212. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10213. int64_t cur_level = maybe_layer->layerId();
  10214. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(window, cur_level)) {
  10215. return at::_ops::stft_center::call(self, n_fft, hop_length, win_length, window, center, pad_mode, normalized, onesided, return_complex, align_to_window);
  10216. }
  10217. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10218. std::optional<Tensor> window_value;
  10219. std::optional<int64_t> window_bdim;
  10220. if (window) {
  10221. std::tie(window_value, window_bdim) = unwrapTensorAtLevel(window.value(), cur_level);
  10222. }
  10223. auto results = batch_rule(self_value, self_bdim, n_fft, hop_length, win_length, window_value, window_bdim, center, pad_mode, normalized, onesided, return_complex, align_to_window);
  10224. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10225. }
  10226. template <typename batch_rule_t, batch_rule_t batch_rule>
  10227. at::Tensor istft_generated_plumbing(const at::Tensor & self, int64_t n_fft, ::std::optional<int64_t> hop_length, ::std::optional<int64_t> win_length, const ::std::optional<at::Tensor> & window, bool center, bool normalized, ::std::optional<bool> onesided, ::std::optional<int64_t> length, bool return_complex) {
  10228. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10229. auto maybe_layer = maybeCurrentDynamicLayer();
  10230. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10231. int64_t cur_level = maybe_layer->layerId();
  10232. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(window, cur_level)) {
  10233. return at::_ops::istft::call(self, n_fft, hop_length, win_length, window, center, normalized, onesided, length, return_complex);
  10234. }
  10235. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10236. std::optional<Tensor> window_value;
  10237. std::optional<int64_t> window_bdim;
  10238. if (window) {
  10239. std::tie(window_value, window_bdim) = unwrapTensorAtLevel(window.value(), cur_level);
  10240. }
  10241. auto results = batch_rule(self_value, self_bdim, n_fft, hop_length, win_length, window_value, window_bdim, center, normalized, onesided, length, return_complex);
  10242. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10243. }
  10244. template <typename batch_rule_t, batch_rule_t batch_rule>
  10245. at::Tensor sum_generated_plumbing(const at::Tensor & self, ::std::optional<at::ScalarType> dtype) {
  10246. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10247. auto maybe_layer = maybeCurrentDynamicLayer();
  10248. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10249. int64_t cur_level = maybe_layer->layerId();
  10250. if (!isBatchedAtLevel(self, cur_level)) {
  10251. return at::_ops::sum::call(self, dtype);
  10252. }
  10253. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10254. auto results = batch_rule(self_value, self_bdim, dtype);
  10255. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10256. }
  10257. template <typename batch_rule_t, batch_rule_t batch_rule>
  10258. at::Tensor sum_dim_IntList_generated_plumbing(const at::Tensor & self, at::OptionalIntArrayRef dim, bool keepdim, ::std::optional<at::ScalarType> dtype) {
  10259. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10260. auto maybe_layer = maybeCurrentDynamicLayer();
  10261. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10262. int64_t cur_level = maybe_layer->layerId();
  10263. if (!isBatchedAtLevel(self, cur_level)) {
  10264. return at::_ops::sum_dim_IntList::call(self, dim, keepdim, dtype);
  10265. }
  10266. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10267. auto results = batch_rule(self_value, self_bdim, dim, keepdim, dtype);
  10268. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10269. }
  10270. template <typename batch_rule_t, batch_rule_t batch_rule>
  10271. at::Tensor sum_dim_DimnameList_generated_plumbing(const at::Tensor & self, at::DimnameList dim, bool keepdim, ::std::optional<at::ScalarType> dtype) {
  10272. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10273. auto maybe_layer = maybeCurrentDynamicLayer();
  10274. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10275. int64_t cur_level = maybe_layer->layerId();
  10276. if (!isBatchedAtLevel(self, cur_level)) {
  10277. return at::_ops::sum_dim_DimnameList::call(self, dim, keepdim, dtype);
  10278. }
  10279. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10280. auto results = batch_rule(self_value, self_bdim, dim, keepdim, dtype);
  10281. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10282. }
  10283. template <typename batch_rule_t, batch_rule_t batch_rule>
  10284. at::Tensor _nested_sum_backward_generated_plumbing(const at::Tensor & grad, const at::Tensor & self, at::OptionalIntArrayRef dim, bool keepdim) {
  10285. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10286. auto maybe_layer = maybeCurrentDynamicLayer();
  10287. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10288. int64_t cur_level = maybe_layer->layerId();
  10289. if (!isBatchedAtLevel(grad, cur_level) && !isBatchedAtLevel(self, cur_level)) {
  10290. return at::_ops::_nested_sum_backward::call(grad, self, dim, keepdim);
  10291. }
  10292. auto [grad_value, grad_bdim] = unwrapTensorAtLevel(grad, cur_level);
  10293. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10294. auto results = batch_rule(grad_value, grad_bdim, self_value, self_bdim, dim, keepdim);
  10295. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10296. }
  10297. template <typename batch_rule_t, batch_rule_t batch_rule>
  10298. at::Tensor nansum_generated_plumbing(const at::Tensor & self, at::OptionalIntArrayRef dim, bool keepdim, ::std::optional<at::ScalarType> dtype) {
  10299. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10300. auto maybe_layer = maybeCurrentDynamicLayer();
  10301. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10302. int64_t cur_level = maybe_layer->layerId();
  10303. if (!isBatchedAtLevel(self, cur_level)) {
  10304. return at::_ops::nansum::call(self, dim, keepdim, dtype);
  10305. }
  10306. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10307. auto results = batch_rule(self_value, self_bdim, dim, keepdim, dtype);
  10308. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10309. }
  10310. template <typename batch_rule_t, batch_rule_t batch_rule>
  10311. at::Tensor hash_tensor_generated_plumbing(const at::Tensor & self, at::IntArrayRef dim, bool keepdim, int64_t mode) {
  10312. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10313. auto maybe_layer = maybeCurrentDynamicLayer();
  10314. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10315. int64_t cur_level = maybe_layer->layerId();
  10316. if (!isBatchedAtLevel(self, cur_level)) {
  10317. return at::_ops::hash_tensor::call(self, dim, keepdim, mode);
  10318. }
  10319. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10320. auto results = batch_rule(self_value, self_bdim, dim, keepdim, mode);
  10321. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10322. }
  10323. template <typename batch_rule_t, batch_rule_t batch_rule>
  10324. at::Tensor sum_to_size_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef size) {
  10325. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10326. auto maybe_layer = maybeCurrentDynamicLayer();
  10327. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10328. int64_t cur_level = maybe_layer->layerId();
  10329. if (!isBatchedAtLevel(self, cur_level)) {
  10330. return at::_ops::sum_to_size::call(self, size);
  10331. }
  10332. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10333. auto results = batch_rule(self_value, self_bdim, size);
  10334. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10335. }
  10336. template <typename batch_rule_t, batch_rule_t batch_rule>
  10337. at::Tensor sqrt_generated_plumbing(const at::Tensor & self) {
  10338. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10339. auto maybe_layer = maybeCurrentDynamicLayer();
  10340. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10341. int64_t cur_level = maybe_layer->layerId();
  10342. if (!isBatchedAtLevel(self, cur_level)) {
  10343. return at::_ops::sqrt::call(self);
  10344. }
  10345. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10346. auto results = batch_rule(self_value, self_bdim);
  10347. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10348. }
  10349. template <typename batch_rule_t, batch_rule_t batch_rule>
  10350. at::Tensor & sqrt__generated_plumbing(at::Tensor & self) {
  10351. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10352. auto maybe_layer = maybeCurrentDynamicLayer();
  10353. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  10354. int64_t cur_level = maybe_layer->layerId();
  10355. if (!isBatchedAtLevel(self, cur_level)) {
  10356. return at::_ops::sqrt_::call(self);
  10357. }
  10358. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10359. batch_rule(self_value, self_bdim);
  10360. return self;
  10361. }
  10362. template <typename batch_rule_t, batch_rule_t batch_rule>
  10363. at::Tensor square_generated_plumbing(const at::Tensor & self) {
  10364. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10365. auto maybe_layer = maybeCurrentDynamicLayer();
  10366. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10367. int64_t cur_level = maybe_layer->layerId();
  10368. if (!isBatchedAtLevel(self, cur_level)) {
  10369. return at::_ops::square::call(self);
  10370. }
  10371. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10372. auto results = batch_rule(self_value, self_bdim);
  10373. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10374. }
  10375. template <typename batch_rule_t, batch_rule_t batch_rule>
  10376. at::Tensor & square__generated_plumbing(at::Tensor & self) {
  10377. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10378. auto maybe_layer = maybeCurrentDynamicLayer();
  10379. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  10380. int64_t cur_level = maybe_layer->layerId();
  10381. if (!isBatchedAtLevel(self, cur_level)) {
  10382. return at::_ops::square_::call(self);
  10383. }
  10384. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10385. batch_rule(self_value, self_bdim);
  10386. return self;
  10387. }
  10388. template <typename batch_rule_t, batch_rule_t batch_rule>
  10389. at::Tensor std_generated_plumbing(const at::Tensor & self, bool unbiased) {
  10390. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10391. auto maybe_layer = maybeCurrentDynamicLayer();
  10392. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10393. int64_t cur_level = maybe_layer->layerId();
  10394. if (!isBatchedAtLevel(self, cur_level)) {
  10395. return at::_ops::std::call(self, unbiased);
  10396. }
  10397. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10398. auto results = batch_rule(self_value, self_bdim, unbiased);
  10399. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10400. }
  10401. template <typename batch_rule_t, batch_rule_t batch_rule>
  10402. at::Tensor std_dim_generated_plumbing(const at::Tensor & self, at::OptionalIntArrayRef dim, bool unbiased, bool keepdim) {
  10403. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10404. auto maybe_layer = maybeCurrentDynamicLayer();
  10405. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10406. int64_t cur_level = maybe_layer->layerId();
  10407. if (!isBatchedAtLevel(self, cur_level)) {
  10408. return at::_ops::std_dim::call(self, dim, unbiased, keepdim);
  10409. }
  10410. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10411. auto results = batch_rule(self_value, self_bdim, dim, unbiased, keepdim);
  10412. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10413. }
  10414. template <typename batch_rule_t, batch_rule_t batch_rule>
  10415. at::Tensor std_correction_generated_plumbing(const at::Tensor & self, at::OptionalIntArrayRef dim, const ::std::optional<at::Scalar> & correction, bool keepdim) {
  10416. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10417. auto maybe_layer = maybeCurrentDynamicLayer();
  10418. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10419. int64_t cur_level = maybe_layer->layerId();
  10420. if (!isBatchedAtLevel(self, cur_level)) {
  10421. return at::_ops::std_correction::call(self, dim, correction, keepdim);
  10422. }
  10423. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10424. auto results = batch_rule(self_value, self_bdim, dim, correction, keepdim);
  10425. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10426. }
  10427. template <typename batch_rule_t, batch_rule_t batch_rule>
  10428. ::std::tuple<at::Tensor,at::Tensor> std_mean_generated_plumbing(const at::Tensor & self, bool unbiased) {
  10429. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10430. auto maybe_layer = maybeCurrentDynamicLayer();
  10431. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10432. int64_t cur_level = maybe_layer->layerId();
  10433. if (!isBatchedAtLevel(self, cur_level)) {
  10434. return at::_ops::std_mean::call(self, unbiased);
  10435. }
  10436. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10437. auto results = batch_rule(self_value, self_bdim, unbiased);
  10438. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  10439. }
  10440. template <typename batch_rule_t, batch_rule_t batch_rule>
  10441. ::std::tuple<at::Tensor,at::Tensor> std_mean_dim_generated_plumbing(const at::Tensor & self, at::OptionalIntArrayRef dim, bool unbiased, bool keepdim) {
  10442. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10443. auto maybe_layer = maybeCurrentDynamicLayer();
  10444. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10445. int64_t cur_level = maybe_layer->layerId();
  10446. if (!isBatchedAtLevel(self, cur_level)) {
  10447. return at::_ops::std_mean_dim::call(self, dim, unbiased, keepdim);
  10448. }
  10449. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10450. auto results = batch_rule(self_value, self_bdim, dim, unbiased, keepdim);
  10451. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  10452. }
  10453. template <typename batch_rule_t, batch_rule_t batch_rule>
  10454. ::std::tuple<at::Tensor,at::Tensor> std_mean_correction_generated_plumbing(const at::Tensor & self, at::OptionalIntArrayRef dim, const ::std::optional<at::Scalar> & correction, bool keepdim) {
  10455. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10456. auto maybe_layer = maybeCurrentDynamicLayer();
  10457. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10458. int64_t cur_level = maybe_layer->layerId();
  10459. if (!isBatchedAtLevel(self, cur_level)) {
  10460. return at::_ops::std_mean_correction::call(self, dim, correction, keepdim);
  10461. }
  10462. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10463. auto results = batch_rule(self_value, self_bdim, dim, correction, keepdim);
  10464. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  10465. }
  10466. template <typename batch_rule_t, batch_rule_t batch_rule>
  10467. ::std::tuple<at::Tensor,at::Tensor> std_mean_names_dim_generated_plumbing(const at::Tensor & self, at::DimnameList dim, bool unbiased, bool keepdim) {
  10468. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10469. auto maybe_layer = maybeCurrentDynamicLayer();
  10470. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10471. int64_t cur_level = maybe_layer->layerId();
  10472. if (!isBatchedAtLevel(self, cur_level)) {
  10473. return at::_ops::std_mean_names_dim::call(self, dim, unbiased, keepdim);
  10474. }
  10475. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10476. auto results = batch_rule(self_value, self_bdim, dim, unbiased, keepdim);
  10477. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  10478. }
  10479. template <typename batch_rule_t, batch_rule_t batch_rule>
  10480. ::std::tuple<at::Tensor,at::Tensor> std_mean_correction_names_generated_plumbing(const at::Tensor & self, at::DimnameList dim, const ::std::optional<at::Scalar> & correction, bool keepdim) {
  10481. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10482. auto maybe_layer = maybeCurrentDynamicLayer();
  10483. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10484. int64_t cur_level = maybe_layer->layerId();
  10485. if (!isBatchedAtLevel(self, cur_level)) {
  10486. return at::_ops::std_mean_correction_names::call(self, dim, correction, keepdim);
  10487. }
  10488. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10489. auto results = batch_rule(self_value, self_bdim, dim, correction, keepdim);
  10490. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  10491. }
  10492. template <typename batch_rule_t, batch_rule_t batch_rule>
  10493. at::Tensor std_names_dim_generated_plumbing(const at::Tensor & self, at::DimnameList dim, bool unbiased, bool keepdim) {
  10494. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10495. auto maybe_layer = maybeCurrentDynamicLayer();
  10496. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10497. int64_t cur_level = maybe_layer->layerId();
  10498. if (!isBatchedAtLevel(self, cur_level)) {
  10499. return at::_ops::std_names_dim::call(self, dim, unbiased, keepdim);
  10500. }
  10501. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10502. auto results = batch_rule(self_value, self_bdim, dim, unbiased, keepdim);
  10503. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10504. }
  10505. template <typename batch_rule_t, batch_rule_t batch_rule>
  10506. at::Tensor std_correction_names_generated_plumbing(const at::Tensor & self, at::DimnameList dim, const ::std::optional<at::Scalar> & correction, bool keepdim) {
  10507. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10508. auto maybe_layer = maybeCurrentDynamicLayer();
  10509. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10510. int64_t cur_level = maybe_layer->layerId();
  10511. if (!isBatchedAtLevel(self, cur_level)) {
  10512. return at::_ops::std_correction_names::call(self, dim, correction, keepdim);
  10513. }
  10514. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10515. auto results = batch_rule(self_value, self_bdim, dim, correction, keepdim);
  10516. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10517. }
  10518. template <typename batch_rule_t, batch_rule_t batch_rule>
  10519. at::Tensor prod_generated_plumbing(const at::Tensor & self, ::std::optional<at::ScalarType> dtype) {
  10520. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10521. auto maybe_layer = maybeCurrentDynamicLayer();
  10522. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10523. int64_t cur_level = maybe_layer->layerId();
  10524. if (!isBatchedAtLevel(self, cur_level)) {
  10525. return at::_ops::prod::call(self, dtype);
  10526. }
  10527. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10528. auto results = batch_rule(self_value, self_bdim, dtype);
  10529. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10530. }
  10531. template <typename batch_rule_t, batch_rule_t batch_rule>
  10532. at::Tensor prod_dim_int_generated_plumbing(const at::Tensor & self, int64_t dim, bool keepdim, ::std::optional<at::ScalarType> dtype) {
  10533. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10534. auto maybe_layer = maybeCurrentDynamicLayer();
  10535. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10536. int64_t cur_level = maybe_layer->layerId();
  10537. if (!isBatchedAtLevel(self, cur_level)) {
  10538. return at::_ops::prod_dim_int::call(self, dim, keepdim, dtype);
  10539. }
  10540. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10541. auto results = batch_rule(self_value, self_bdim, dim, keepdim, dtype);
  10542. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10543. }
  10544. template <typename batch_rule_t, batch_rule_t batch_rule>
  10545. at::Tensor prod_dim_Dimname_generated_plumbing(const at::Tensor & self, at::Dimname dim, bool keepdim, ::std::optional<at::ScalarType> dtype) {
  10546. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10547. auto maybe_layer = maybeCurrentDynamicLayer();
  10548. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10549. int64_t cur_level = maybe_layer->layerId();
  10550. if (!isBatchedAtLevel(self, cur_level)) {
  10551. return at::_ops::prod_dim_Dimname::call(self, dim, keepdim, dtype);
  10552. }
  10553. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10554. auto results = batch_rule(self_value, self_bdim, dim, keepdim, dtype);
  10555. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10556. }
  10557. template <typename batch_rule_t, batch_rule_t batch_rule>
  10558. at::Tensor t_generated_plumbing(const at::Tensor & self) {
  10559. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10560. auto maybe_layer = maybeCurrentDynamicLayer();
  10561. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10562. int64_t cur_level = maybe_layer->layerId();
  10563. if (!isBatchedAtLevel(self, cur_level)) {
  10564. return at::_ops::t::call(self);
  10565. }
  10566. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10567. auto results = batch_rule(self_value, self_bdim);
  10568. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10569. }
  10570. template <typename batch_rule_t, batch_rule_t batch_rule>
  10571. at::Tensor tan_generated_plumbing(const at::Tensor & self) {
  10572. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10573. auto maybe_layer = maybeCurrentDynamicLayer();
  10574. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10575. int64_t cur_level = maybe_layer->layerId();
  10576. if (!isBatchedAtLevel(self, cur_level)) {
  10577. return at::_ops::tan::call(self);
  10578. }
  10579. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10580. auto results = batch_rule(self_value, self_bdim);
  10581. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10582. }
  10583. template <typename batch_rule_t, batch_rule_t batch_rule>
  10584. at::Tensor & tan__generated_plumbing(at::Tensor & self) {
  10585. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10586. auto maybe_layer = maybeCurrentDynamicLayer();
  10587. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  10588. int64_t cur_level = maybe_layer->layerId();
  10589. if (!isBatchedAtLevel(self, cur_level)) {
  10590. return at::_ops::tan_::call(self);
  10591. }
  10592. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10593. batch_rule(self_value, self_bdim);
  10594. return self;
  10595. }
  10596. template <typename batch_rule_t, batch_rule_t batch_rule>
  10597. at::Tensor tanh_generated_plumbing(const at::Tensor & self) {
  10598. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10599. auto maybe_layer = maybeCurrentDynamicLayer();
  10600. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10601. int64_t cur_level = maybe_layer->layerId();
  10602. if (!isBatchedAtLevel(self, cur_level)) {
  10603. return at::_ops::tanh::call(self);
  10604. }
  10605. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10606. auto results = batch_rule(self_value, self_bdim);
  10607. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10608. }
  10609. template <typename batch_rule_t, batch_rule_t batch_rule>
  10610. at::Tensor & tanh__generated_plumbing(at::Tensor & self) {
  10611. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10612. auto maybe_layer = maybeCurrentDynamicLayer();
  10613. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  10614. int64_t cur_level = maybe_layer->layerId();
  10615. if (!isBatchedAtLevel(self, cur_level)) {
  10616. return at::_ops::tanh_::call(self);
  10617. }
  10618. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10619. batch_rule(self_value, self_bdim);
  10620. return self;
  10621. }
  10622. template <typename batch_rule_t, batch_rule_t batch_rule>
  10623. at::Tensor tensordot_generated_plumbing(const at::Tensor & self, const at::Tensor & other, at::IntArrayRef dims_self, at::IntArrayRef dims_other) {
  10624. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10625. auto maybe_layer = maybeCurrentDynamicLayer();
  10626. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10627. int64_t cur_level = maybe_layer->layerId();
  10628. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  10629. return at::_ops::tensordot::call(self, other, dims_self, dims_other);
  10630. }
  10631. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10632. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  10633. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim, dims_self, dims_other);
  10634. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10635. }
  10636. template <typename batch_rule_t, batch_rule_t batch_rule>
  10637. at::Tensor threshold_generated_plumbing(const at::Tensor & self, const at::Scalar & threshold, const at::Scalar & value) {
  10638. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10639. auto maybe_layer = maybeCurrentDynamicLayer();
  10640. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10641. int64_t cur_level = maybe_layer->layerId();
  10642. if (!isBatchedAtLevel(self, cur_level)) {
  10643. return at::_ops::threshold::call(self, threshold, value);
  10644. }
  10645. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10646. auto results = batch_rule(self_value, self_bdim, threshold, value);
  10647. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10648. }
  10649. template <typename batch_rule_t, batch_rule_t batch_rule>
  10650. at::Tensor & threshold__generated_plumbing(at::Tensor & self, const at::Scalar & threshold, const at::Scalar & value) {
  10651. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10652. auto maybe_layer = maybeCurrentDynamicLayer();
  10653. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  10654. int64_t cur_level = maybe_layer->layerId();
  10655. if (!isBatchedAtLevel(self, cur_level)) {
  10656. return at::_ops::threshold_::call(self, threshold, value);
  10657. }
  10658. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10659. batch_rule(self_value, self_bdim, threshold, value);
  10660. return self;
  10661. }
  10662. template <typename batch_rule_t, batch_rule_t batch_rule>
  10663. at::Tensor threshold_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self, const at::Scalar & threshold) {
  10664. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10665. auto maybe_layer = maybeCurrentDynamicLayer();
  10666. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10667. int64_t cur_level = maybe_layer->layerId();
  10668. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level)) {
  10669. return at::_ops::threshold_backward::call(grad_output, self, threshold);
  10670. }
  10671. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  10672. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10673. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim, threshold);
  10674. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10675. }
  10676. template <typename batch_rule_t, batch_rule_t batch_rule>
  10677. at::Tensor tile_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef dims) {
  10678. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10679. auto maybe_layer = maybeCurrentDynamicLayer();
  10680. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10681. int64_t cur_level = maybe_layer->layerId();
  10682. if (!isBatchedAtLevel(self, cur_level)) {
  10683. return at::_ops::tile::call(self, dims);
  10684. }
  10685. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10686. auto results = batch_rule(self_value, self_bdim, dims);
  10687. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10688. }
  10689. template <typename batch_rule_t, batch_rule_t batch_rule>
  10690. at::Tensor transpose_int_generated_plumbing(const at::Tensor & self, int64_t dim0, int64_t dim1) {
  10691. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10692. auto maybe_layer = maybeCurrentDynamicLayer();
  10693. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10694. int64_t cur_level = maybe_layer->layerId();
  10695. if (!isBatchedAtLevel(self, cur_level)) {
  10696. return at::_ops::transpose_int::call(self, dim0, dim1);
  10697. }
  10698. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10699. auto results = batch_rule(self_value, self_bdim, dim0, dim1);
  10700. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10701. }
  10702. template <typename batch_rule_t, batch_rule_t batch_rule>
  10703. at::Tensor transpose_Dimname_generated_plumbing(const at::Tensor & self, at::Dimname dim0, at::Dimname dim1) {
  10704. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10705. auto maybe_layer = maybeCurrentDynamicLayer();
  10706. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10707. int64_t cur_level = maybe_layer->layerId();
  10708. if (!isBatchedAtLevel(self, cur_level)) {
  10709. return at::_ops::transpose_Dimname::call(self, dim0, dim1);
  10710. }
  10711. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10712. auto results = batch_rule(self_value, self_bdim, dim0, dim1);
  10713. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10714. }
  10715. template <typename batch_rule_t, batch_rule_t batch_rule>
  10716. at::Tensor _mkldnn_transpose_generated_plumbing(const at::Tensor & self, int64_t dim0, int64_t dim1) {
  10717. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10718. auto maybe_layer = maybeCurrentDynamicLayer();
  10719. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10720. int64_t cur_level = maybe_layer->layerId();
  10721. if (!isBatchedAtLevel(self, cur_level)) {
  10722. return at::_ops::_mkldnn_transpose::call(self, dim0, dim1);
  10723. }
  10724. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10725. auto results = batch_rule(self_value, self_bdim, dim0, dim1);
  10726. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10727. }
  10728. template <typename batch_rule_t, batch_rule_t batch_rule>
  10729. at::Tensor & _mkldnn_transpose__generated_plumbing(at::Tensor & self, int64_t dim0, int64_t dim1) {
  10730. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10731. auto maybe_layer = maybeCurrentDynamicLayer();
  10732. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  10733. int64_t cur_level = maybe_layer->layerId();
  10734. if (!isBatchedAtLevel(self, cur_level)) {
  10735. return at::_ops::_mkldnn_transpose_::call(self, dim0, dim1);
  10736. }
  10737. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10738. batch_rule(self_value, self_bdim, dim0, dim1);
  10739. return self;
  10740. }
  10741. template <typename batch_rule_t, batch_rule_t batch_rule>
  10742. at::Tensor one_hot_generated_plumbing(const at::Tensor & self, int64_t num_classes) {
  10743. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10744. auto maybe_layer = maybeCurrentDynamicLayer();
  10745. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10746. int64_t cur_level = maybe_layer->layerId();
  10747. if (!isBatchedAtLevel(self, cur_level)) {
  10748. return at::_ops::one_hot::call(self, num_classes);
  10749. }
  10750. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10751. auto results = batch_rule(self_value, self_bdim, num_classes);
  10752. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10753. }
  10754. template <typename batch_rule_t, batch_rule_t batch_rule>
  10755. at::Tensor flip_generated_plumbing(const at::Tensor & self, at::IntArrayRef dims) {
  10756. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10757. auto maybe_layer = maybeCurrentDynamicLayer();
  10758. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10759. int64_t cur_level = maybe_layer->layerId();
  10760. if (!isBatchedAtLevel(self, cur_level)) {
  10761. return at::_ops::flip::call(self, dims);
  10762. }
  10763. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10764. auto results = batch_rule(self_value, self_bdim, dims);
  10765. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10766. }
  10767. template <typename batch_rule_t, batch_rule_t batch_rule>
  10768. at::Tensor fliplr_generated_plumbing(const at::Tensor & self) {
  10769. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10770. auto maybe_layer = maybeCurrentDynamicLayer();
  10771. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10772. int64_t cur_level = maybe_layer->layerId();
  10773. if (!isBatchedAtLevel(self, cur_level)) {
  10774. return at::_ops::fliplr::call(self);
  10775. }
  10776. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10777. auto results = batch_rule(self_value, self_bdim);
  10778. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10779. }
  10780. template <typename batch_rule_t, batch_rule_t batch_rule>
  10781. at::Tensor flipud_generated_plumbing(const at::Tensor & self) {
  10782. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10783. auto maybe_layer = maybeCurrentDynamicLayer();
  10784. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10785. int64_t cur_level = maybe_layer->layerId();
  10786. if (!isBatchedAtLevel(self, cur_level)) {
  10787. return at::_ops::flipud::call(self);
  10788. }
  10789. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10790. auto results = batch_rule(self_value, self_bdim);
  10791. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10792. }
  10793. template <typename batch_rule_t, batch_rule_t batch_rule>
  10794. at::Tensor roll_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef shifts, at::IntArrayRef dims) {
  10795. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10796. auto maybe_layer = maybeCurrentDynamicLayer();
  10797. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10798. int64_t cur_level = maybe_layer->layerId();
  10799. if (!isBatchedAtLevel(self, cur_level)) {
  10800. return at::_ops::roll::call(self, shifts, dims);
  10801. }
  10802. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10803. auto results = batch_rule(self_value, self_bdim, shifts, dims);
  10804. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10805. }
  10806. template <typename batch_rule_t, batch_rule_t batch_rule>
  10807. at::Tensor rot90_generated_plumbing(const at::Tensor & self, int64_t k, at::IntArrayRef dims) {
  10808. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10809. auto maybe_layer = maybeCurrentDynamicLayer();
  10810. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10811. int64_t cur_level = maybe_layer->layerId();
  10812. if (!isBatchedAtLevel(self, cur_level)) {
  10813. return at::_ops::rot90::call(self, k, dims);
  10814. }
  10815. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10816. auto results = batch_rule(self_value, self_bdim, k, dims);
  10817. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10818. }
  10819. template <typename batch_rule_t, batch_rule_t batch_rule>
  10820. at::Tensor trapezoid_x_generated_plumbing(const at::Tensor & y, const at::Tensor & x, int64_t dim) {
  10821. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10822. auto maybe_layer = maybeCurrentDynamicLayer();
  10823. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10824. int64_t cur_level = maybe_layer->layerId();
  10825. if (!isBatchedAtLevel(y, cur_level) && !isBatchedAtLevel(x, cur_level)) {
  10826. return at::_ops::trapezoid_x::call(y, x, dim);
  10827. }
  10828. auto [y_value, y_bdim] = unwrapTensorAtLevel(y, cur_level);
  10829. auto [x_value, x_bdim] = unwrapTensorAtLevel(x, cur_level);
  10830. auto results = batch_rule(y_value, y_bdim, x_value, x_bdim, dim);
  10831. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10832. }
  10833. template <typename batch_rule_t, batch_rule_t batch_rule>
  10834. at::Tensor trapezoid_dx_generated_plumbing(const at::Tensor & y, const at::Scalar & dx, int64_t dim) {
  10835. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10836. auto maybe_layer = maybeCurrentDynamicLayer();
  10837. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10838. int64_t cur_level = maybe_layer->layerId();
  10839. if (!isBatchedAtLevel(y, cur_level)) {
  10840. return at::_ops::trapezoid_dx::call(y, dx, dim);
  10841. }
  10842. auto [y_value, y_bdim] = unwrapTensorAtLevel(y, cur_level);
  10843. auto results = batch_rule(y_value, y_bdim, dx, dim);
  10844. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10845. }
  10846. template <typename batch_rule_t, batch_rule_t batch_rule>
  10847. at::Tensor trapz_x_generated_plumbing(const at::Tensor & y, const at::Tensor & x, int64_t dim) {
  10848. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10849. auto maybe_layer = maybeCurrentDynamicLayer();
  10850. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10851. int64_t cur_level = maybe_layer->layerId();
  10852. if (!isBatchedAtLevel(y, cur_level) && !isBatchedAtLevel(x, cur_level)) {
  10853. return at::_ops::trapz_x::call(y, x, dim);
  10854. }
  10855. auto [y_value, y_bdim] = unwrapTensorAtLevel(y, cur_level);
  10856. auto [x_value, x_bdim] = unwrapTensorAtLevel(x, cur_level);
  10857. auto results = batch_rule(y_value, y_bdim, x_value, x_bdim, dim);
  10858. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10859. }
  10860. template <typename batch_rule_t, batch_rule_t batch_rule>
  10861. at::Tensor trapz_dx_generated_plumbing(const at::Tensor & y, double dx, int64_t dim) {
  10862. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10863. auto maybe_layer = maybeCurrentDynamicLayer();
  10864. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10865. int64_t cur_level = maybe_layer->layerId();
  10866. if (!isBatchedAtLevel(y, cur_level)) {
  10867. return at::_ops::trapz_dx::call(y, dx, dim);
  10868. }
  10869. auto [y_value, y_bdim] = unwrapTensorAtLevel(y, cur_level);
  10870. auto results = batch_rule(y_value, y_bdim, dx, dim);
  10871. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10872. }
  10873. template <typename batch_rule_t, batch_rule_t batch_rule>
  10874. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> _transform_bias_rescale_qkv_generated_plumbing(const at::Tensor & qkv, const at::Tensor & qkv_bias, int64_t num_heads) {
  10875. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10876. auto maybe_layer = maybeCurrentDynamicLayer();
  10877. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10878. int64_t cur_level = maybe_layer->layerId();
  10879. if (!isBatchedAtLevel(qkv, cur_level) && !isBatchedAtLevel(qkv_bias, cur_level)) {
  10880. return at::_ops::_transform_bias_rescale_qkv::call(qkv, qkv_bias, num_heads);
  10881. }
  10882. auto [qkv_value, qkv_bdim] = unwrapTensorAtLevel(qkv, cur_level);
  10883. auto [qkv_bias_value, qkv_bias_bdim] = unwrapTensorAtLevel(qkv_bias, cur_level);
  10884. auto results = batch_rule(qkv_value, qkv_bdim, qkv_bias_value, qkv_bias_bdim, num_heads);
  10885. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  10886. }
  10887. template <typename batch_rule_t, batch_rule_t batch_rule>
  10888. at::Tensor _nested_tensor_from_mask_generated_plumbing(const at::Tensor & t, const at::Tensor & mask, bool mask_check) {
  10889. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10890. auto maybe_layer = maybeCurrentDynamicLayer();
  10891. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10892. int64_t cur_level = maybe_layer->layerId();
  10893. if (!isBatchedAtLevel(t, cur_level) && !isBatchedAtLevel(mask, cur_level)) {
  10894. return at::_ops::_nested_tensor_from_mask::call(t, mask, mask_check);
  10895. }
  10896. auto [t_value, t_bdim] = unwrapTensorAtLevel(t, cur_level);
  10897. auto [mask_value, mask_bdim] = unwrapTensorAtLevel(mask, cur_level);
  10898. auto results = batch_rule(t_value, t_bdim, mask_value, mask_bdim, mask_check);
  10899. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10900. }
  10901. template <typename batch_rule_t, batch_rule_t batch_rule>
  10902. at::Tensor _nested_from_padded_generated_plumbing(const at::Tensor & padded, const at::Tensor & cpu_nested_shape_example, bool fuse_transform_0213) {
  10903. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10904. auto maybe_layer = maybeCurrentDynamicLayer();
  10905. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10906. int64_t cur_level = maybe_layer->layerId();
  10907. if (!isBatchedAtLevel(padded, cur_level) && !isBatchedAtLevel(cpu_nested_shape_example, cur_level)) {
  10908. return at::_ops::_nested_from_padded::call(padded, cpu_nested_shape_example, fuse_transform_0213);
  10909. }
  10910. auto [padded_value, padded_bdim] = unwrapTensorAtLevel(padded, cur_level);
  10911. auto [cpu_nested_shape_example_value, cpu_nested_shape_example_bdim] = unwrapTensorAtLevel(cpu_nested_shape_example, cur_level);
  10912. auto results = batch_rule(padded_value, padded_bdim, cpu_nested_shape_example_value, cpu_nested_shape_example_bdim, fuse_transform_0213);
  10913. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10914. }
  10915. template <typename batch_rule_t, batch_rule_t batch_rule>
  10916. at::Tensor _nested_tensor_size_generated_plumbing(const at::Tensor & self) {
  10917. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10918. auto maybe_layer = maybeCurrentDynamicLayer();
  10919. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10920. int64_t cur_level = maybe_layer->layerId();
  10921. if (!isBatchedAtLevel(self, cur_level)) {
  10922. return at::_ops::_nested_tensor_size::call(self);
  10923. }
  10924. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10925. auto results = batch_rule(self_value, self_bdim);
  10926. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10927. }
  10928. template <typename batch_rule_t, batch_rule_t batch_rule>
  10929. at::Tensor _nested_tensor_strides_generated_plumbing(const at::Tensor & self) {
  10930. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10931. auto maybe_layer = maybeCurrentDynamicLayer();
  10932. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10933. int64_t cur_level = maybe_layer->layerId();
  10934. if (!isBatchedAtLevel(self, cur_level)) {
  10935. return at::_ops::_nested_tensor_strides::call(self);
  10936. }
  10937. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10938. auto results = batch_rule(self_value, self_bdim);
  10939. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10940. }
  10941. template <typename batch_rule_t, batch_rule_t batch_rule>
  10942. at::Tensor _nested_tensor_storage_offsets_generated_plumbing(const at::Tensor & self) {
  10943. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10944. auto maybe_layer = maybeCurrentDynamicLayer();
  10945. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10946. int64_t cur_level = maybe_layer->layerId();
  10947. if (!isBatchedAtLevel(self, cur_level)) {
  10948. return at::_ops::_nested_tensor_storage_offsets::call(self);
  10949. }
  10950. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10951. auto results = batch_rule(self_value, self_bdim);
  10952. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10953. }
  10954. template <typename batch_rule_t, batch_rule_t batch_rule>
  10955. at::Tensor _nested_from_padded_and_nested_example_generated_plumbing(const at::Tensor & padded, const at::Tensor & nt_example) {
  10956. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10957. auto maybe_layer = maybeCurrentDynamicLayer();
  10958. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10959. int64_t cur_level = maybe_layer->layerId();
  10960. if (!isBatchedAtLevel(padded, cur_level) && !isBatchedAtLevel(nt_example, cur_level)) {
  10961. return at::_ops::_nested_from_padded_and_nested_example::call(padded, nt_example);
  10962. }
  10963. auto [padded_value, padded_bdim] = unwrapTensorAtLevel(padded, cur_level);
  10964. auto [nt_example_value, nt_example_bdim] = unwrapTensorAtLevel(nt_example, cur_level);
  10965. auto results = batch_rule(padded_value, padded_bdim, nt_example_value, nt_example_bdim);
  10966. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10967. }
  10968. template <typename batch_rule_t, batch_rule_t batch_rule>
  10969. at::Tensor _nested_view_from_buffer_generated_plumbing(const at::Tensor & self, const at::Tensor & nested_size, const at::Tensor & nested_strides, const at::Tensor & offsets) {
  10970. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10971. auto maybe_layer = maybeCurrentDynamicLayer();
  10972. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10973. int64_t cur_level = maybe_layer->layerId();
  10974. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(nested_size, cur_level) && !isBatchedAtLevel(nested_strides, cur_level) && !isBatchedAtLevel(offsets, cur_level)) {
  10975. return at::_ops::_nested_view_from_buffer::call(self, nested_size, nested_strides, offsets);
  10976. }
  10977. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10978. auto [nested_size_value, nested_size_bdim] = unwrapTensorAtLevel(nested_size, cur_level);
  10979. auto [nested_strides_value, nested_strides_bdim] = unwrapTensorAtLevel(nested_strides, cur_level);
  10980. auto [offsets_value, offsets_bdim] = unwrapTensorAtLevel(offsets, cur_level);
  10981. auto results = batch_rule(self_value, self_bdim, nested_size_value, nested_size_bdim, nested_strides_value, nested_strides_bdim, offsets_value, offsets_bdim);
  10982. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10983. }
  10984. template <typename batch_rule_t, batch_rule_t batch_rule>
  10985. at::Tensor _nested_view_from_buffer_copy_generated_plumbing(const at::Tensor & self, const at::Tensor & nested_size, const at::Tensor & nested_strides, const at::Tensor & offsets) {
  10986. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  10987. auto maybe_layer = maybeCurrentDynamicLayer();
  10988. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  10989. int64_t cur_level = maybe_layer->layerId();
  10990. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(nested_size, cur_level) && !isBatchedAtLevel(nested_strides, cur_level) && !isBatchedAtLevel(offsets, cur_level)) {
  10991. return at::_ops::_nested_view_from_buffer_copy::call(self, nested_size, nested_strides, offsets);
  10992. }
  10993. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  10994. auto [nested_size_value, nested_size_bdim] = unwrapTensorAtLevel(nested_size, cur_level);
  10995. auto [nested_strides_value, nested_strides_bdim] = unwrapTensorAtLevel(nested_strides, cur_level);
  10996. auto [offsets_value, offsets_bdim] = unwrapTensorAtLevel(offsets, cur_level);
  10997. auto results = batch_rule(self_value, self_bdim, nested_size_value, nested_size_bdim, nested_strides_value, nested_strides_bdim, offsets_value, offsets_bdim);
  10998. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  10999. }
  11000. template <typename batch_rule_t, batch_rule_t batch_rule>
  11001. at::Tensor _nested_view_from_jagged_generated_plumbing(const at::Tensor & self, const at::Tensor & offsets, const at::Tensor & dummy, const ::std::optional<at::Tensor> & lengths, int64_t ragged_idx, const ::std::optional<at::Tensor> & min_seqlen, const ::std::optional<at::Tensor> & max_seqlen) {
  11002. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11003. auto maybe_layer = maybeCurrentDynamicLayer();
  11004. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11005. int64_t cur_level = maybe_layer->layerId();
  11006. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(offsets, cur_level) && !isBatchedAtLevel(dummy, cur_level) && !isBatchedAtLevel(lengths, cur_level) && !isBatchedAtLevel(min_seqlen, cur_level) && !isBatchedAtLevel(max_seqlen, cur_level)) {
  11007. return at::_ops::_nested_view_from_jagged::call(self, offsets, dummy, lengths, ragged_idx, min_seqlen, max_seqlen);
  11008. }
  11009. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11010. auto [offsets_value, offsets_bdim] = unwrapTensorAtLevel(offsets, cur_level);
  11011. auto [dummy_value, dummy_bdim] = unwrapTensorAtLevel(dummy, cur_level);
  11012. std::optional<Tensor> lengths_value;
  11013. std::optional<int64_t> lengths_bdim;
  11014. if (lengths) {
  11015. std::tie(lengths_value, lengths_bdim) = unwrapTensorAtLevel(lengths.value(), cur_level);
  11016. }
  11017. std::optional<Tensor> min_seqlen_value;
  11018. std::optional<int64_t> min_seqlen_bdim;
  11019. if (min_seqlen) {
  11020. std::tie(min_seqlen_value, min_seqlen_bdim) = unwrapTensorAtLevel(min_seqlen.value(), cur_level);
  11021. }
  11022. std::optional<Tensor> max_seqlen_value;
  11023. std::optional<int64_t> max_seqlen_bdim;
  11024. if (max_seqlen) {
  11025. std::tie(max_seqlen_value, max_seqlen_bdim) = unwrapTensorAtLevel(max_seqlen.value(), cur_level);
  11026. }
  11027. auto results = batch_rule(self_value, self_bdim, offsets_value, offsets_bdim, dummy_value, dummy_bdim, lengths_value, lengths_bdim, ragged_idx, min_seqlen_value, min_seqlen_bdim, max_seqlen_value, max_seqlen_bdim);
  11028. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11029. }
  11030. template <typename batch_rule_t, batch_rule_t batch_rule>
  11031. at::Tensor _nested_view_from_jagged_copy_generated_plumbing(const at::Tensor & self, const at::Tensor & offsets, const at::Tensor & dummy, const ::std::optional<at::Tensor> & lengths, int64_t ragged_idx, const ::std::optional<at::Tensor> & min_seqlen, const ::std::optional<at::Tensor> & max_seqlen) {
  11032. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11033. auto maybe_layer = maybeCurrentDynamicLayer();
  11034. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11035. int64_t cur_level = maybe_layer->layerId();
  11036. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(offsets, cur_level) && !isBatchedAtLevel(dummy, cur_level) && !isBatchedAtLevel(lengths, cur_level) && !isBatchedAtLevel(min_seqlen, cur_level) && !isBatchedAtLevel(max_seqlen, cur_level)) {
  11037. return at::_ops::_nested_view_from_jagged_copy::call(self, offsets, dummy, lengths, ragged_idx, min_seqlen, max_seqlen);
  11038. }
  11039. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11040. auto [offsets_value, offsets_bdim] = unwrapTensorAtLevel(offsets, cur_level);
  11041. auto [dummy_value, dummy_bdim] = unwrapTensorAtLevel(dummy, cur_level);
  11042. std::optional<Tensor> lengths_value;
  11043. std::optional<int64_t> lengths_bdim;
  11044. if (lengths) {
  11045. std::tie(lengths_value, lengths_bdim) = unwrapTensorAtLevel(lengths.value(), cur_level);
  11046. }
  11047. std::optional<Tensor> min_seqlen_value;
  11048. std::optional<int64_t> min_seqlen_bdim;
  11049. if (min_seqlen) {
  11050. std::tie(min_seqlen_value, min_seqlen_bdim) = unwrapTensorAtLevel(min_seqlen.value(), cur_level);
  11051. }
  11052. std::optional<Tensor> max_seqlen_value;
  11053. std::optional<int64_t> max_seqlen_bdim;
  11054. if (max_seqlen) {
  11055. std::tie(max_seqlen_value, max_seqlen_bdim) = unwrapTensorAtLevel(max_seqlen.value(), cur_level);
  11056. }
  11057. auto results = batch_rule(self_value, self_bdim, offsets_value, offsets_bdim, dummy_value, dummy_bdim, lengths_value, lengths_bdim, ragged_idx, min_seqlen_value, min_seqlen_bdim, max_seqlen_value, max_seqlen_bdim);
  11058. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11059. }
  11060. template <typename batch_rule_t, batch_rule_t batch_rule>
  11061. at::Tensor _nested_get_values_generated_plumbing(const at::Tensor & self) {
  11062. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11063. auto maybe_layer = maybeCurrentDynamicLayer();
  11064. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11065. int64_t cur_level = maybe_layer->layerId();
  11066. if (!isBatchedAtLevel(self, cur_level)) {
  11067. return at::_ops::_nested_get_values::call(self);
  11068. }
  11069. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11070. auto results = batch_rule(self_value, self_bdim);
  11071. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11072. }
  11073. template <typename batch_rule_t, batch_rule_t batch_rule>
  11074. at::Tensor _nested_get_values_copy_generated_plumbing(const at::Tensor & self) {
  11075. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11076. auto maybe_layer = maybeCurrentDynamicLayer();
  11077. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11078. int64_t cur_level = maybe_layer->layerId();
  11079. if (!isBatchedAtLevel(self, cur_level)) {
  11080. return at::_ops::_nested_get_values_copy::call(self);
  11081. }
  11082. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11083. auto results = batch_rule(self_value, self_bdim);
  11084. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11085. }
  11086. template <typename batch_rule_t, batch_rule_t batch_rule>
  11087. at::Tensor _nested_get_offsets_generated_plumbing(const at::Tensor & self) {
  11088. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11089. auto maybe_layer = maybeCurrentDynamicLayer();
  11090. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11091. int64_t cur_level = maybe_layer->layerId();
  11092. if (!isBatchedAtLevel(self, cur_level)) {
  11093. return at::_ops::_nested_get_offsets::call(self);
  11094. }
  11095. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11096. auto results = batch_rule(self_value, self_bdim);
  11097. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11098. }
  11099. template <typename batch_rule_t, batch_rule_t batch_rule>
  11100. at::Tensor _nested_get_lengths_generated_plumbing(const at::Tensor & self) {
  11101. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11102. auto maybe_layer = maybeCurrentDynamicLayer();
  11103. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11104. int64_t cur_level = maybe_layer->layerId();
  11105. if (!isBatchedAtLevel(self, cur_level)) {
  11106. return at::_ops::_nested_get_lengths::call(self);
  11107. }
  11108. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11109. auto results = batch_rule(self_value, self_bdim);
  11110. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11111. }
  11112. template <typename batch_rule_t, batch_rule_t batch_rule>
  11113. at::Tensor _nested_get_min_seqlen_generated_plumbing(const at::Tensor & self) {
  11114. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11115. auto maybe_layer = maybeCurrentDynamicLayer();
  11116. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11117. int64_t cur_level = maybe_layer->layerId();
  11118. if (!isBatchedAtLevel(self, cur_level)) {
  11119. return at::_ops::_nested_get_min_seqlen::call(self);
  11120. }
  11121. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11122. auto results = batch_rule(self_value, self_bdim);
  11123. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11124. }
  11125. template <typename batch_rule_t, batch_rule_t batch_rule>
  11126. at::Tensor _nested_get_max_seqlen_generated_plumbing(const at::Tensor & self) {
  11127. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11128. auto maybe_layer = maybeCurrentDynamicLayer();
  11129. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11130. int64_t cur_level = maybe_layer->layerId();
  11131. if (!isBatchedAtLevel(self, cur_level)) {
  11132. return at::_ops::_nested_get_max_seqlen::call(self);
  11133. }
  11134. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11135. auto results = batch_rule(self_value, self_bdim);
  11136. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11137. }
  11138. template <typename batch_rule_t, batch_rule_t batch_rule>
  11139. at::Tensor _nested_get_jagged_dummy_generated_plumbing(const at::Tensor & any) {
  11140. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11141. auto maybe_layer = maybeCurrentDynamicLayer();
  11142. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11143. int64_t cur_level = maybe_layer->layerId();
  11144. if (!isBatchedAtLevel(any, cur_level)) {
  11145. return at::_ops::_nested_get_jagged_dummy::call(any);
  11146. }
  11147. auto [any_value, any_bdim] = unwrapTensorAtLevel(any, cur_level);
  11148. auto results = batch_rule(any_value, any_bdim);
  11149. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11150. }
  11151. template <typename batch_rule_t, batch_rule_t batch_rule>
  11152. ::std::tuple<at::Tensor,at::Tensor> _nested_compute_contiguous_strides_offsets_generated_plumbing(const at::Tensor & nested_size) {
  11153. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11154. auto maybe_layer = maybeCurrentDynamicLayer();
  11155. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11156. int64_t cur_level = maybe_layer->layerId();
  11157. if (!isBatchedAtLevel(nested_size, cur_level)) {
  11158. return at::_ops::_nested_compute_contiguous_strides_offsets::call(nested_size);
  11159. }
  11160. auto [nested_size_value, nested_size_bdim] = unwrapTensorAtLevel(nested_size, cur_level);
  11161. auto results = batch_rule(nested_size_value, nested_size_bdim);
  11162. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  11163. }
  11164. template <typename batch_rule_t, batch_rule_t batch_rule>
  11165. at::Tensor _trilinear_generated_plumbing(const at::Tensor & i1, const at::Tensor & i2, const at::Tensor & i3, at::IntArrayRef expand1, at::IntArrayRef expand2, at::IntArrayRef expand3, at::IntArrayRef sumdim, int64_t unroll_dim) {
  11166. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11167. auto maybe_layer = maybeCurrentDynamicLayer();
  11168. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11169. int64_t cur_level = maybe_layer->layerId();
  11170. if (!isBatchedAtLevel(i1, cur_level) && !isBatchedAtLevel(i2, cur_level) && !isBatchedAtLevel(i3, cur_level)) {
  11171. return at::_ops::_trilinear::call(i1, i2, i3, expand1, expand2, expand3, sumdim, unroll_dim);
  11172. }
  11173. auto [i1_value, i1_bdim] = unwrapTensorAtLevel(i1, cur_level);
  11174. auto [i2_value, i2_bdim] = unwrapTensorAtLevel(i2, cur_level);
  11175. auto [i3_value, i3_bdim] = unwrapTensorAtLevel(i3, cur_level);
  11176. auto results = batch_rule(i1_value, i1_bdim, i2_value, i2_bdim, i3_value, i3_bdim, expand1, expand2, expand3, sumdim, unroll_dim);
  11177. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11178. }
  11179. template <typename batch_rule_t, batch_rule_t batch_rule>
  11180. at::Tensor triplet_margin_loss_generated_plumbing(const at::Tensor & anchor, const at::Tensor & positive, const at::Tensor & negative, double margin, double p, double eps, bool swap, int64_t reduction) {
  11181. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11182. auto maybe_layer = maybeCurrentDynamicLayer();
  11183. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11184. int64_t cur_level = maybe_layer->layerId();
  11185. if (!isBatchedAtLevel(anchor, cur_level) && !isBatchedAtLevel(positive, cur_level) && !isBatchedAtLevel(negative, cur_level)) {
  11186. return at::_ops::triplet_margin_loss::call(anchor, positive, negative, margin, p, eps, swap, reduction);
  11187. }
  11188. auto [anchor_value, anchor_bdim] = unwrapTensorAtLevel(anchor, cur_level);
  11189. auto [positive_value, positive_bdim] = unwrapTensorAtLevel(positive, cur_level);
  11190. auto [negative_value, negative_bdim] = unwrapTensorAtLevel(negative, cur_level);
  11191. auto results = batch_rule(anchor_value, anchor_bdim, positive_value, positive_bdim, negative_value, negative_bdim, margin, p, eps, swap, reduction);
  11192. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11193. }
  11194. template <typename batch_rule_t, batch_rule_t batch_rule>
  11195. at::Tensor trunc_generated_plumbing(const at::Tensor & self) {
  11196. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11197. auto maybe_layer = maybeCurrentDynamicLayer();
  11198. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11199. int64_t cur_level = maybe_layer->layerId();
  11200. if (!isBatchedAtLevel(self, cur_level)) {
  11201. return at::_ops::trunc::call(self);
  11202. }
  11203. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11204. auto results = batch_rule(self_value, self_bdim);
  11205. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11206. }
  11207. template <typename batch_rule_t, batch_rule_t batch_rule>
  11208. at::Tensor & trunc__generated_plumbing(at::Tensor & self) {
  11209. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11210. auto maybe_layer = maybeCurrentDynamicLayer();
  11211. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  11212. int64_t cur_level = maybe_layer->layerId();
  11213. if (!isBatchedAtLevel(self, cur_level)) {
  11214. return at::_ops::trunc_::call(self);
  11215. }
  11216. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11217. batch_rule(self_value, self_bdim);
  11218. return self;
  11219. }
  11220. template <typename batch_rule_t, batch_rule_t batch_rule>
  11221. at::Tensor fix_generated_plumbing(const at::Tensor & self) {
  11222. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11223. auto maybe_layer = maybeCurrentDynamicLayer();
  11224. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11225. int64_t cur_level = maybe_layer->layerId();
  11226. if (!isBatchedAtLevel(self, cur_level)) {
  11227. return at::_ops::fix::call(self);
  11228. }
  11229. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11230. auto results = batch_rule(self_value, self_bdim);
  11231. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11232. }
  11233. template <typename batch_rule_t, batch_rule_t batch_rule>
  11234. at::Tensor & fix__generated_plumbing(at::Tensor & self) {
  11235. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11236. auto maybe_layer = maybeCurrentDynamicLayer();
  11237. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  11238. int64_t cur_level = maybe_layer->layerId();
  11239. if (!isBatchedAtLevel(self, cur_level)) {
  11240. return at::_ops::fix_::call(self);
  11241. }
  11242. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11243. batch_rule(self_value, self_bdim);
  11244. return self;
  11245. }
  11246. template <typename batch_rule_t, batch_rule_t batch_rule>
  11247. at::Tensor type_as_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  11248. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11249. auto maybe_layer = maybeCurrentDynamicLayer();
  11250. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11251. int64_t cur_level = maybe_layer->layerId();
  11252. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  11253. return at::_ops::type_as::call(self, other);
  11254. }
  11255. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11256. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  11257. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  11258. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11259. }
  11260. template <typename batch_rule_t, batch_rule_t batch_rule>
  11261. ::std::tuple<at::Tensor,at::Tensor> _unique_generated_plumbing(const at::Tensor & self, bool sorted, bool return_inverse) {
  11262. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11263. auto maybe_layer = maybeCurrentDynamicLayer();
  11264. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11265. int64_t cur_level = maybe_layer->layerId();
  11266. if (!isBatchedAtLevel(self, cur_level)) {
  11267. return at::_ops::_unique::call(self, sorted, return_inverse);
  11268. }
  11269. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11270. auto results = batch_rule(self_value, self_bdim, sorted, return_inverse);
  11271. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  11272. }
  11273. template <typename batch_rule_t, batch_rule_t batch_rule>
  11274. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> unique_dim_generated_plumbing(const at::Tensor & self, int64_t dim, bool sorted, bool return_inverse, bool return_counts) {
  11275. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11276. auto maybe_layer = maybeCurrentDynamicLayer();
  11277. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11278. int64_t cur_level = maybe_layer->layerId();
  11279. if (!isBatchedAtLevel(self, cur_level)) {
  11280. return at::_ops::unique_dim::call(self, dim, sorted, return_inverse, return_counts);
  11281. }
  11282. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11283. auto results = batch_rule(self_value, self_bdim, dim, sorted, return_inverse, return_counts);
  11284. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  11285. }
  11286. template <typename batch_rule_t, batch_rule_t batch_rule>
  11287. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> unique_consecutive_generated_plumbing(const at::Tensor & self, bool return_inverse, bool return_counts, ::std::optional<int64_t> dim) {
  11288. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11289. auto maybe_layer = maybeCurrentDynamicLayer();
  11290. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11291. int64_t cur_level = maybe_layer->layerId();
  11292. if (!isBatchedAtLevel(self, cur_level)) {
  11293. return at::_ops::unique_consecutive::call(self, return_inverse, return_counts, dim);
  11294. }
  11295. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11296. auto results = batch_rule(self_value, self_bdim, return_inverse, return_counts, dim);
  11297. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  11298. }
  11299. template <typename batch_rule_t, batch_rule_t batch_rule>
  11300. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> unique_dim_consecutive_generated_plumbing(const at::Tensor & self, int64_t dim, bool return_inverse, bool return_counts) {
  11301. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11302. auto maybe_layer = maybeCurrentDynamicLayer();
  11303. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11304. int64_t cur_level = maybe_layer->layerId();
  11305. if (!isBatchedAtLevel(self, cur_level)) {
  11306. return at::_ops::unique_dim_consecutive::call(self, dim, return_inverse, return_counts);
  11307. }
  11308. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11309. auto results = batch_rule(self_value, self_bdim, dim, return_inverse, return_counts);
  11310. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  11311. }
  11312. template <typename batch_rule_t, batch_rule_t batch_rule>
  11313. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> _unique2_generated_plumbing(const at::Tensor & self, bool sorted, bool return_inverse, bool return_counts) {
  11314. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11315. auto maybe_layer = maybeCurrentDynamicLayer();
  11316. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11317. int64_t cur_level = maybe_layer->layerId();
  11318. if (!isBatchedAtLevel(self, cur_level)) {
  11319. return at::_ops::_unique2::call(self, sorted, return_inverse, return_counts);
  11320. }
  11321. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11322. auto results = batch_rule(self_value, self_bdim, sorted, return_inverse, return_counts);
  11323. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  11324. }
  11325. template <typename batch_rule_t, batch_rule_t batch_rule>
  11326. at::Tensor _unsafe_view_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef size) {
  11327. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11328. auto maybe_layer = maybeCurrentDynamicLayer();
  11329. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11330. int64_t cur_level = maybe_layer->layerId();
  11331. if (!isBatchedAtLevel(self, cur_level)) {
  11332. return at::_ops::_unsafe_view::call(self, size);
  11333. }
  11334. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11335. auto results = batch_rule(self_value, self_bdim, size);
  11336. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11337. }
  11338. template <typename batch_rule_t, batch_rule_t batch_rule>
  11339. at::Tensor unsqueeze_generated_plumbing(const at::Tensor & self, int64_t dim) {
  11340. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11341. auto maybe_layer = maybeCurrentDynamicLayer();
  11342. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11343. int64_t cur_level = maybe_layer->layerId();
  11344. if (!isBatchedAtLevel(self, cur_level)) {
  11345. return at::_ops::unsqueeze::call(self, dim);
  11346. }
  11347. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11348. auto results = batch_rule(self_value, self_bdim, dim);
  11349. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11350. }
  11351. template <typename batch_rule_t, batch_rule_t batch_rule>
  11352. at::Tensor vander_generated_plumbing(const at::Tensor & x, ::std::optional<int64_t> N, bool increasing) {
  11353. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11354. auto maybe_layer = maybeCurrentDynamicLayer();
  11355. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11356. int64_t cur_level = maybe_layer->layerId();
  11357. if (!isBatchedAtLevel(x, cur_level)) {
  11358. return at::_ops::vander::call(x, N, increasing);
  11359. }
  11360. auto [x_value, x_bdim] = unwrapTensorAtLevel(x, cur_level);
  11361. auto results = batch_rule(x_value, x_bdim, N, increasing);
  11362. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11363. }
  11364. template <typename batch_rule_t, batch_rule_t batch_rule>
  11365. at::Tensor var_generated_plumbing(const at::Tensor & self, bool unbiased) {
  11366. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11367. auto maybe_layer = maybeCurrentDynamicLayer();
  11368. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11369. int64_t cur_level = maybe_layer->layerId();
  11370. if (!isBatchedAtLevel(self, cur_level)) {
  11371. return at::_ops::var::call(self, unbiased);
  11372. }
  11373. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11374. auto results = batch_rule(self_value, self_bdim, unbiased);
  11375. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11376. }
  11377. template <typename batch_rule_t, batch_rule_t batch_rule>
  11378. at::Tensor var_dim_generated_plumbing(const at::Tensor & self, at::OptionalIntArrayRef dim, bool unbiased, bool keepdim) {
  11379. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11380. auto maybe_layer = maybeCurrentDynamicLayer();
  11381. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11382. int64_t cur_level = maybe_layer->layerId();
  11383. if (!isBatchedAtLevel(self, cur_level)) {
  11384. return at::_ops::var_dim::call(self, dim, unbiased, keepdim);
  11385. }
  11386. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11387. auto results = batch_rule(self_value, self_bdim, dim, unbiased, keepdim);
  11388. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11389. }
  11390. template <typename batch_rule_t, batch_rule_t batch_rule>
  11391. at::Tensor var_correction_generated_plumbing(const at::Tensor & self, at::OptionalIntArrayRef dim, const ::std::optional<at::Scalar> & correction, bool keepdim) {
  11392. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11393. auto maybe_layer = maybeCurrentDynamicLayer();
  11394. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11395. int64_t cur_level = maybe_layer->layerId();
  11396. if (!isBatchedAtLevel(self, cur_level)) {
  11397. return at::_ops::var_correction::call(self, dim, correction, keepdim);
  11398. }
  11399. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11400. auto results = batch_rule(self_value, self_bdim, dim, correction, keepdim);
  11401. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11402. }
  11403. template <typename batch_rule_t, batch_rule_t batch_rule>
  11404. at::Tensor var_names_dim_generated_plumbing(const at::Tensor & self, at::DimnameList dim, bool unbiased, bool keepdim) {
  11405. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11406. auto maybe_layer = maybeCurrentDynamicLayer();
  11407. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11408. int64_t cur_level = maybe_layer->layerId();
  11409. if (!isBatchedAtLevel(self, cur_level)) {
  11410. return at::_ops::var_names_dim::call(self, dim, unbiased, keepdim);
  11411. }
  11412. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11413. auto results = batch_rule(self_value, self_bdim, dim, unbiased, keepdim);
  11414. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11415. }
  11416. template <typename batch_rule_t, batch_rule_t batch_rule>
  11417. at::Tensor var_correction_names_generated_plumbing(const at::Tensor & self, at::DimnameList dim, const ::std::optional<at::Scalar> & correction, bool keepdim) {
  11418. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11419. auto maybe_layer = maybeCurrentDynamicLayer();
  11420. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11421. int64_t cur_level = maybe_layer->layerId();
  11422. if (!isBatchedAtLevel(self, cur_level)) {
  11423. return at::_ops::var_correction_names::call(self, dim, correction, keepdim);
  11424. }
  11425. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11426. auto results = batch_rule(self_value, self_bdim, dim, correction, keepdim);
  11427. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11428. }
  11429. template <typename batch_rule_t, batch_rule_t batch_rule>
  11430. ::std::tuple<at::Tensor,at::Tensor> var_mean_generated_plumbing(const at::Tensor & self, bool unbiased) {
  11431. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11432. auto maybe_layer = maybeCurrentDynamicLayer();
  11433. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11434. int64_t cur_level = maybe_layer->layerId();
  11435. if (!isBatchedAtLevel(self, cur_level)) {
  11436. return at::_ops::var_mean::call(self, unbiased);
  11437. }
  11438. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11439. auto results = batch_rule(self_value, self_bdim, unbiased);
  11440. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  11441. }
  11442. template <typename batch_rule_t, batch_rule_t batch_rule>
  11443. ::std::tuple<at::Tensor,at::Tensor> var_mean_dim_generated_plumbing(const at::Tensor & self, at::OptionalIntArrayRef dim, bool unbiased, bool keepdim) {
  11444. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11445. auto maybe_layer = maybeCurrentDynamicLayer();
  11446. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11447. int64_t cur_level = maybe_layer->layerId();
  11448. if (!isBatchedAtLevel(self, cur_level)) {
  11449. return at::_ops::var_mean_dim::call(self, dim, unbiased, keepdim);
  11450. }
  11451. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11452. auto results = batch_rule(self_value, self_bdim, dim, unbiased, keepdim);
  11453. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  11454. }
  11455. template <typename batch_rule_t, batch_rule_t batch_rule>
  11456. ::std::tuple<at::Tensor,at::Tensor> var_mean_correction_generated_plumbing(const at::Tensor & self, at::OptionalIntArrayRef dim, const ::std::optional<at::Scalar> & correction, bool keepdim) {
  11457. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11458. auto maybe_layer = maybeCurrentDynamicLayer();
  11459. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11460. int64_t cur_level = maybe_layer->layerId();
  11461. if (!isBatchedAtLevel(self, cur_level)) {
  11462. return at::_ops::var_mean_correction::call(self, dim, correction, keepdim);
  11463. }
  11464. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11465. auto results = batch_rule(self_value, self_bdim, dim, correction, keepdim);
  11466. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  11467. }
  11468. template <typename batch_rule_t, batch_rule_t batch_rule>
  11469. ::std::tuple<at::Tensor,at::Tensor> var_mean_names_dim_generated_plumbing(const at::Tensor & self, at::DimnameList dim, bool unbiased, bool keepdim) {
  11470. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11471. auto maybe_layer = maybeCurrentDynamicLayer();
  11472. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11473. int64_t cur_level = maybe_layer->layerId();
  11474. if (!isBatchedAtLevel(self, cur_level)) {
  11475. return at::_ops::var_mean_names_dim::call(self, dim, unbiased, keepdim);
  11476. }
  11477. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11478. auto results = batch_rule(self_value, self_bdim, dim, unbiased, keepdim);
  11479. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  11480. }
  11481. template <typename batch_rule_t, batch_rule_t batch_rule>
  11482. ::std::tuple<at::Tensor,at::Tensor> var_mean_correction_names_generated_plumbing(const at::Tensor & self, at::DimnameList dim, const ::std::optional<at::Scalar> & correction, bool keepdim) {
  11483. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11484. auto maybe_layer = maybeCurrentDynamicLayer();
  11485. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11486. int64_t cur_level = maybe_layer->layerId();
  11487. if (!isBatchedAtLevel(self, cur_level)) {
  11488. return at::_ops::var_mean_correction_names::call(self, dim, correction, keepdim);
  11489. }
  11490. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11491. auto results = batch_rule(self_value, self_bdim, dim, correction, keepdim);
  11492. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  11493. }
  11494. template <typename batch_rule_t, batch_rule_t batch_rule>
  11495. at::Tensor view_as_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  11496. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11497. auto maybe_layer = maybeCurrentDynamicLayer();
  11498. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11499. int64_t cur_level = maybe_layer->layerId();
  11500. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  11501. return at::_ops::view_as::call(self, other);
  11502. }
  11503. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11504. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  11505. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  11506. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11507. }
  11508. template <typename batch_rule_t, batch_rule_t batch_rule>
  11509. at::Tensor where_self_generated_plumbing(const at::Tensor & condition, const at::Tensor & self, const at::Tensor & other) {
  11510. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11511. auto maybe_layer = maybeCurrentDynamicLayer();
  11512. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11513. int64_t cur_level = maybe_layer->layerId();
  11514. if (!isBatchedAtLevel(condition, cur_level) && !isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  11515. return at::_ops::where_self::call(condition, self, other);
  11516. }
  11517. auto [condition_value, condition_bdim] = unwrapTensorAtLevel(condition, cur_level);
  11518. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11519. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  11520. auto results = batch_rule(condition_value, condition_bdim, self_value, self_bdim, other_value, other_bdim);
  11521. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11522. }
  11523. template <typename batch_rule_t, batch_rule_t batch_rule>
  11524. at::Tensor where_ScalarSelf_generated_plumbing(const at::Tensor & condition, const at::Scalar & self, const at::Tensor & other) {
  11525. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11526. auto maybe_layer = maybeCurrentDynamicLayer();
  11527. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11528. int64_t cur_level = maybe_layer->layerId();
  11529. if (!isBatchedAtLevel(condition, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  11530. return at::_ops::where_ScalarSelf::call(condition, self, other);
  11531. }
  11532. auto [condition_value, condition_bdim] = unwrapTensorAtLevel(condition, cur_level);
  11533. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  11534. auto results = batch_rule(condition_value, condition_bdim, self, other_value, other_bdim);
  11535. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11536. }
  11537. template <typename batch_rule_t, batch_rule_t batch_rule>
  11538. at::Tensor where_ScalarOther_generated_plumbing(const at::Tensor & condition, const at::Tensor & self, const at::Scalar & other) {
  11539. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11540. auto maybe_layer = maybeCurrentDynamicLayer();
  11541. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11542. int64_t cur_level = maybe_layer->layerId();
  11543. if (!isBatchedAtLevel(condition, cur_level) && !isBatchedAtLevel(self, cur_level)) {
  11544. return at::_ops::where_ScalarOther::call(condition, self, other);
  11545. }
  11546. auto [condition_value, condition_bdim] = unwrapTensorAtLevel(condition, cur_level);
  11547. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11548. auto results = batch_rule(condition_value, condition_bdim, self_value, self_bdim, other);
  11549. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11550. }
  11551. template <typename batch_rule_t, batch_rule_t batch_rule>
  11552. at::Tensor where_Scalar_generated_plumbing(const at::Tensor & condition, const at::Scalar & self, const at::Scalar & other) {
  11553. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11554. auto maybe_layer = maybeCurrentDynamicLayer();
  11555. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11556. int64_t cur_level = maybe_layer->layerId();
  11557. if (!isBatchedAtLevel(condition, cur_level)) {
  11558. return at::_ops::where_Scalar::call(condition, self, other);
  11559. }
  11560. auto [condition_value, condition_bdim] = unwrapTensorAtLevel(condition, cur_level);
  11561. auto results = batch_rule(condition_value, condition_bdim, self, other);
  11562. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11563. }
  11564. template <typename batch_rule_t, batch_rule_t batch_rule>
  11565. ::std::vector<at::Tensor> where_generated_plumbing(const at::Tensor & condition) {
  11566. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11567. auto maybe_layer = maybeCurrentDynamicLayer();
  11568. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11569. int64_t cur_level = maybe_layer->layerId();
  11570. if (!isBatchedAtLevel(condition, cur_level)) {
  11571. return at::_ops::where::call(condition);
  11572. }
  11573. auto [condition_value, condition_bdim] = unwrapTensorAtLevel(condition, cur_level);
  11574. auto results = batch_rule(condition_value, condition_bdim);
  11575. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  11576. }
  11577. template <typename batch_rule_t, batch_rule_t batch_rule>
  11578. at::Tensor norm_except_dim_generated_plumbing(const at::Tensor & v, int64_t pow, int64_t dim) {
  11579. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11580. auto maybe_layer = maybeCurrentDynamicLayer();
  11581. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11582. int64_t cur_level = maybe_layer->layerId();
  11583. if (!isBatchedAtLevel(v, cur_level)) {
  11584. return at::_ops::norm_except_dim::call(v, pow, dim);
  11585. }
  11586. auto [v_value, v_bdim] = unwrapTensorAtLevel(v, cur_level);
  11587. auto results = batch_rule(v_value, v_bdim, pow, dim);
  11588. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11589. }
  11590. template <typename batch_rule_t, batch_rule_t batch_rule>
  11591. at::Tensor _weight_norm_generated_plumbing(const at::Tensor & v, const at::Tensor & g, int64_t dim) {
  11592. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11593. auto maybe_layer = maybeCurrentDynamicLayer();
  11594. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11595. int64_t cur_level = maybe_layer->layerId();
  11596. if (!isBatchedAtLevel(v, cur_level) && !isBatchedAtLevel(g, cur_level)) {
  11597. return at::_ops::_weight_norm::call(v, g, dim);
  11598. }
  11599. auto [v_value, v_bdim] = unwrapTensorAtLevel(v, cur_level);
  11600. auto [g_value, g_bdim] = unwrapTensorAtLevel(g, cur_level);
  11601. auto results = batch_rule(v_value, v_bdim, g_value, g_bdim, dim);
  11602. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11603. }
  11604. template <typename batch_rule_t, batch_rule_t batch_rule>
  11605. ::std::tuple<at::Tensor,at::Tensor> _weight_norm_interface_generated_plumbing(const at::Tensor & v, const at::Tensor & g, int64_t dim) {
  11606. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11607. auto maybe_layer = maybeCurrentDynamicLayer();
  11608. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11609. int64_t cur_level = maybe_layer->layerId();
  11610. if (!isBatchedAtLevel(v, cur_level) && !isBatchedAtLevel(g, cur_level)) {
  11611. return at::_ops::_weight_norm_interface::call(v, g, dim);
  11612. }
  11613. auto [v_value, v_bdim] = unwrapTensorAtLevel(v, cur_level);
  11614. auto [g_value, g_bdim] = unwrapTensorAtLevel(g, cur_level);
  11615. auto results = batch_rule(v_value, v_bdim, g_value, g_bdim, dim);
  11616. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  11617. }
  11618. template <typename batch_rule_t, batch_rule_t batch_rule>
  11619. ::std::tuple<at::Tensor,at::Tensor> _weight_norm_interface_backward_generated_plumbing(const at::Tensor & grad_w, const at::Tensor & saved_v, const at::Tensor & saved_g, const at::Tensor & saved_norms, int64_t dim) {
  11620. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11621. auto maybe_layer = maybeCurrentDynamicLayer();
  11622. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11623. int64_t cur_level = maybe_layer->layerId();
  11624. if (!isBatchedAtLevel(grad_w, cur_level) && !isBatchedAtLevel(saved_v, cur_level) && !isBatchedAtLevel(saved_g, cur_level) && !isBatchedAtLevel(saved_norms, cur_level)) {
  11625. return at::_ops::_weight_norm_interface_backward::call(grad_w, saved_v, saved_g, saved_norms, dim);
  11626. }
  11627. auto [grad_w_value, grad_w_bdim] = unwrapTensorAtLevel(grad_w, cur_level);
  11628. auto [saved_v_value, saved_v_bdim] = unwrapTensorAtLevel(saved_v, cur_level);
  11629. auto [saved_g_value, saved_g_bdim] = unwrapTensorAtLevel(saved_g, cur_level);
  11630. auto [saved_norms_value, saved_norms_bdim] = unwrapTensorAtLevel(saved_norms, cur_level);
  11631. auto results = batch_rule(grad_w_value, grad_w_bdim, saved_v_value, saved_v_bdim, saved_g_value, saved_g_bdim, saved_norms_value, saved_norms_bdim, dim);
  11632. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  11633. }
  11634. template <typename batch_rule_t, batch_rule_t batch_rule>
  11635. ::std::tuple<at::Tensor,at::Tensor> _weight_norm_differentiable_backward_generated_plumbing(const at::Tensor & grad_w, const at::Tensor & saved_v, const at::Tensor & saved_g, const at::Tensor & saved_norms, int64_t dim) {
  11636. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11637. auto maybe_layer = maybeCurrentDynamicLayer();
  11638. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11639. int64_t cur_level = maybe_layer->layerId();
  11640. if (!isBatchedAtLevel(grad_w, cur_level) && !isBatchedAtLevel(saved_v, cur_level) && !isBatchedAtLevel(saved_g, cur_level) && !isBatchedAtLevel(saved_norms, cur_level)) {
  11641. return at::_ops::_weight_norm_differentiable_backward::call(grad_w, saved_v, saved_g, saved_norms, dim);
  11642. }
  11643. auto [grad_w_value, grad_w_bdim] = unwrapTensorAtLevel(grad_w, cur_level);
  11644. auto [saved_v_value, saved_v_bdim] = unwrapTensorAtLevel(saved_v, cur_level);
  11645. auto [saved_g_value, saved_g_bdim] = unwrapTensorAtLevel(saved_g, cur_level);
  11646. auto [saved_norms_value, saved_norms_bdim] = unwrapTensorAtLevel(saved_norms, cur_level);
  11647. auto results = batch_rule(grad_w_value, grad_w_bdim, saved_v_value, saved_v_bdim, saved_g_value, saved_g_bdim, saved_norms_value, saved_norms_bdim, dim);
  11648. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  11649. }
  11650. template <typename batch_rule_t, batch_rule_t batch_rule>
  11651. at::Tensor zeros_like_generated_plumbing(const at::Tensor & self, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory, ::std::optional<at::MemoryFormat> memory_format) {
  11652. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11653. auto maybe_layer = maybeCurrentDynamicLayer();
  11654. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11655. int64_t cur_level = maybe_layer->layerId();
  11656. if (!isBatchedAtLevel(self, cur_level)) {
  11657. return at::_ops::zeros_like::call(self, dtype, layout, device, pin_memory, memory_format);
  11658. }
  11659. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11660. auto results = batch_rule(self_value, self_bdim, dtype, layout, device, pin_memory, memory_format);
  11661. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11662. }
  11663. template <typename batch_rule_t, batch_rule_t batch_rule>
  11664. at::Tensor _standard_gamma_grad_generated_plumbing(const at::Tensor & self, const at::Tensor & output) {
  11665. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11666. auto maybe_layer = maybeCurrentDynamicLayer();
  11667. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11668. int64_t cur_level = maybe_layer->layerId();
  11669. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(output, cur_level)) {
  11670. return at::_ops::_standard_gamma_grad::call(self, output);
  11671. }
  11672. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11673. auto [output_value, output_bdim] = unwrapTensorAtLevel(output, cur_level);
  11674. auto results = batch_rule(self_value, self_bdim, output_value, output_bdim);
  11675. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11676. }
  11677. template <typename batch_rule_t, batch_rule_t batch_rule>
  11678. at::Tensor _standard_gamma_generated_plumbing(const at::Tensor & self, ::std::optional<at::Generator> generator) {
  11679. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11680. auto maybe_layer = maybeCurrentDynamicLayer();
  11681. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11682. int64_t cur_level = maybe_layer->layerId();
  11683. if (!isBatchedAtLevel(self, cur_level)) {
  11684. return at::_ops::_standard_gamma::call(self, generator);
  11685. }
  11686. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11687. auto results = batch_rule(self_value, self_bdim, generator);
  11688. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11689. }
  11690. template <typename batch_rule_t, batch_rule_t batch_rule>
  11691. at::Tensor _dirichlet_grad_generated_plumbing(const at::Tensor & x, const at::Tensor & alpha, const at::Tensor & total) {
  11692. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11693. auto maybe_layer = maybeCurrentDynamicLayer();
  11694. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11695. int64_t cur_level = maybe_layer->layerId();
  11696. if (!isBatchedAtLevel(x, cur_level) && !isBatchedAtLevel(alpha, cur_level) && !isBatchedAtLevel(total, cur_level)) {
  11697. return at::_ops::_dirichlet_grad::call(x, alpha, total);
  11698. }
  11699. auto [x_value, x_bdim] = unwrapTensorAtLevel(x, cur_level);
  11700. auto [alpha_value, alpha_bdim] = unwrapTensorAtLevel(alpha, cur_level);
  11701. auto [total_value, total_bdim] = unwrapTensorAtLevel(total, cur_level);
  11702. auto results = batch_rule(x_value, x_bdim, alpha_value, alpha_bdim, total_value, total_bdim);
  11703. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11704. }
  11705. template <typename batch_rule_t, batch_rule_t batch_rule>
  11706. at::Tensor _sample_dirichlet_generated_plumbing(const at::Tensor & self, ::std::optional<at::Generator> generator) {
  11707. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11708. auto maybe_layer = maybeCurrentDynamicLayer();
  11709. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11710. int64_t cur_level = maybe_layer->layerId();
  11711. if (!isBatchedAtLevel(self, cur_level)) {
  11712. return at::_ops::_sample_dirichlet::call(self, generator);
  11713. }
  11714. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11715. auto results = batch_rule(self_value, self_bdim, generator);
  11716. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11717. }
  11718. template <typename batch_rule_t, batch_rule_t batch_rule>
  11719. at::Tensor poisson_generated_plumbing(const at::Tensor & self, ::std::optional<at::Generator> generator) {
  11720. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11721. auto maybe_layer = maybeCurrentDynamicLayer();
  11722. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11723. int64_t cur_level = maybe_layer->layerId();
  11724. if (!isBatchedAtLevel(self, cur_level)) {
  11725. return at::_ops::poisson::call(self, generator);
  11726. }
  11727. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11728. auto results = batch_rule(self_value, self_bdim, generator);
  11729. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11730. }
  11731. template <typename batch_rule_t, batch_rule_t batch_rule>
  11732. at::Tensor binomial_generated_plumbing(const at::Tensor & count, const at::Tensor & prob, ::std::optional<at::Generator> generator) {
  11733. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11734. auto maybe_layer = maybeCurrentDynamicLayer();
  11735. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11736. int64_t cur_level = maybe_layer->layerId();
  11737. if (!isBatchedAtLevel(count, cur_level) && !isBatchedAtLevel(prob, cur_level)) {
  11738. return at::_ops::binomial::call(count, prob, generator);
  11739. }
  11740. auto [count_value, count_bdim] = unwrapTensorAtLevel(count, cur_level);
  11741. auto [prob_value, prob_bdim] = unwrapTensorAtLevel(prob, cur_level);
  11742. auto results = batch_rule(count_value, count_bdim, prob_value, prob_bdim, generator);
  11743. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11744. }
  11745. template <typename batch_rule_t, batch_rule_t batch_rule>
  11746. at::Tensor native_norm_generated_plumbing(const at::Tensor & self, const at::Scalar & p) {
  11747. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11748. auto maybe_layer = maybeCurrentDynamicLayer();
  11749. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11750. int64_t cur_level = maybe_layer->layerId();
  11751. if (!isBatchedAtLevel(self, cur_level)) {
  11752. return at::_ops::native_norm::call(self, p);
  11753. }
  11754. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11755. auto results = batch_rule(self_value, self_bdim, p);
  11756. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11757. }
  11758. template <typename batch_rule_t, batch_rule_t batch_rule>
  11759. at::Tensor native_norm_ScalarOpt_dim_dtype_generated_plumbing(const at::Tensor & self, const ::std::optional<at::Scalar> & p, at::IntArrayRef dim, bool keepdim, ::std::optional<at::ScalarType> dtype) {
  11760. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11761. auto maybe_layer = maybeCurrentDynamicLayer();
  11762. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11763. int64_t cur_level = maybe_layer->layerId();
  11764. if (!isBatchedAtLevel(self, cur_level)) {
  11765. return at::_ops::native_norm_ScalarOpt_dim_dtype::call(self, p, dim, keepdim, dtype);
  11766. }
  11767. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11768. auto results = batch_rule(self_value, self_bdim, p, dim, keepdim, dtype);
  11769. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11770. }
  11771. template <typename batch_rule_t, batch_rule_t batch_rule>
  11772. ::std::tuple<at::Tensor,at::Tensor,at::Tensor,at::Tensor> _batch_norm_no_update_generated_plumbing(const at::Tensor & input, const ::std::optional<at::Tensor> & weight, const ::std::optional<at::Tensor> & bias, const ::std::optional<at::Tensor> & running_mean, const ::std::optional<at::Tensor> & running_var, double momentum, double eps) {
  11773. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11774. auto maybe_layer = maybeCurrentDynamicLayer();
  11775. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11776. int64_t cur_level = maybe_layer->layerId();
  11777. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level) && !isBatchedAtLevel(running_mean, cur_level) && !isBatchedAtLevel(running_var, cur_level)) {
  11778. return at::_ops::_batch_norm_no_update::call(input, weight, bias, running_mean, running_var, momentum, eps);
  11779. }
  11780. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  11781. std::optional<Tensor> weight_value;
  11782. std::optional<int64_t> weight_bdim;
  11783. if (weight) {
  11784. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  11785. }
  11786. std::optional<Tensor> bias_value;
  11787. std::optional<int64_t> bias_bdim;
  11788. if (bias) {
  11789. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  11790. }
  11791. std::optional<Tensor> running_mean_value;
  11792. std::optional<int64_t> running_mean_bdim;
  11793. if (running_mean) {
  11794. std::tie(running_mean_value, running_mean_bdim) = unwrapTensorAtLevel(running_mean.value(), cur_level);
  11795. }
  11796. std::optional<Tensor> running_var_value;
  11797. std::optional<int64_t> running_var_bdim;
  11798. if (running_var) {
  11799. std::tie(running_var_value, running_var_bdim) = unwrapTensorAtLevel(running_var.value(), cur_level);
  11800. }
  11801. auto results = batch_rule(input_value, input_bdim, weight_value, weight_bdim, bias_value, bias_bdim, running_mean_value, running_mean_bdim, running_var_value, running_var_bdim, momentum, eps);
  11802. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level), makeBatched(std::get<6>(results), std::get<7>(results), cur_level));
  11803. }
  11804. template <typename batch_rule_t, batch_rule_t batch_rule>
  11805. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> batch_norm_backward_generated_plumbing(const at::Tensor & grad_out, const at::Tensor & input, const at::Tensor & weight, const ::std::optional<at::Tensor> & running_mean, const ::std::optional<at::Tensor> & running_var, const ::std::optional<at::Tensor> & save_mean, const ::std::optional<at::Tensor> & save_var, bool update, double eps, ::std::array<bool,3> output_mask, const at::Tensor & reserve) {
  11806. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11807. auto maybe_layer = maybeCurrentDynamicLayer();
  11808. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11809. int64_t cur_level = maybe_layer->layerId();
  11810. if (!isBatchedAtLevel(grad_out, cur_level) && !isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(running_mean, cur_level) && !isBatchedAtLevel(running_var, cur_level) && !isBatchedAtLevel(save_mean, cur_level) && !isBatchedAtLevel(save_var, cur_level) && !isBatchedAtLevel(reserve, cur_level)) {
  11811. return at::_ops::batch_norm_backward::call(grad_out, input, weight, running_mean, running_var, save_mean, save_var, update, eps, output_mask, reserve);
  11812. }
  11813. auto [grad_out_value, grad_out_bdim] = unwrapTensorAtLevel(grad_out, cur_level);
  11814. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  11815. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  11816. auto [reserve_value, reserve_bdim] = unwrapTensorAtLevel(reserve, cur_level);
  11817. std::optional<Tensor> running_mean_value;
  11818. std::optional<int64_t> running_mean_bdim;
  11819. if (running_mean) {
  11820. std::tie(running_mean_value, running_mean_bdim) = unwrapTensorAtLevel(running_mean.value(), cur_level);
  11821. }
  11822. std::optional<Tensor> running_var_value;
  11823. std::optional<int64_t> running_var_bdim;
  11824. if (running_var) {
  11825. std::tie(running_var_value, running_var_bdim) = unwrapTensorAtLevel(running_var.value(), cur_level);
  11826. }
  11827. std::optional<Tensor> save_mean_value;
  11828. std::optional<int64_t> save_mean_bdim;
  11829. if (save_mean) {
  11830. std::tie(save_mean_value, save_mean_bdim) = unwrapTensorAtLevel(save_mean.value(), cur_level);
  11831. }
  11832. std::optional<Tensor> save_var_value;
  11833. std::optional<int64_t> save_var_bdim;
  11834. if (save_var) {
  11835. std::tie(save_var_value, save_var_bdim) = unwrapTensorAtLevel(save_var.value(), cur_level);
  11836. }
  11837. auto results = batch_rule(grad_out_value, grad_out_bdim, input_value, input_bdim, weight_value, weight_bdim, running_mean_value, running_mean_bdim, running_var_value, running_var_bdim, save_mean_value, save_mean_bdim, save_var_value, save_var_bdim, update, eps, output_mask, reserve_value, reserve_bdim);
  11838. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  11839. }
  11840. template <typename batch_rule_t, batch_rule_t batch_rule>
  11841. at::Tensor _sparse_sum_generated_plumbing(const at::Tensor & self) {
  11842. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11843. auto maybe_layer = maybeCurrentDynamicLayer();
  11844. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11845. int64_t cur_level = maybe_layer->layerId();
  11846. if (!isBatchedAtLevel(self, cur_level)) {
  11847. return at::_ops::_sparse_sum::call(self);
  11848. }
  11849. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11850. auto results = batch_rule(self_value, self_bdim);
  11851. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11852. }
  11853. template <typename batch_rule_t, batch_rule_t batch_rule>
  11854. at::Tensor _sparse_sum_dtype_generated_plumbing(const at::Tensor & self, at::ScalarType dtype) {
  11855. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11856. auto maybe_layer = maybeCurrentDynamicLayer();
  11857. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11858. int64_t cur_level = maybe_layer->layerId();
  11859. if (!isBatchedAtLevel(self, cur_level)) {
  11860. return at::_ops::_sparse_sum_dtype::call(self, dtype);
  11861. }
  11862. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11863. auto results = batch_rule(self_value, self_bdim, dtype);
  11864. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11865. }
  11866. template <typename batch_rule_t, batch_rule_t batch_rule>
  11867. at::Tensor _sparse_sum_dim_generated_plumbing(const at::Tensor & self, at::IntArrayRef dim) {
  11868. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11869. auto maybe_layer = maybeCurrentDynamicLayer();
  11870. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11871. int64_t cur_level = maybe_layer->layerId();
  11872. if (!isBatchedAtLevel(self, cur_level)) {
  11873. return at::_ops::_sparse_sum_dim::call(self, dim);
  11874. }
  11875. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11876. auto results = batch_rule(self_value, self_bdim, dim);
  11877. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11878. }
  11879. template <typename batch_rule_t, batch_rule_t batch_rule>
  11880. at::Tensor _sparse_sum_dim_dtype_generated_plumbing(const at::Tensor & self, at::IntArrayRef dim, at::ScalarType dtype) {
  11881. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11882. auto maybe_layer = maybeCurrentDynamicLayer();
  11883. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11884. int64_t cur_level = maybe_layer->layerId();
  11885. if (!isBatchedAtLevel(self, cur_level)) {
  11886. return at::_ops::_sparse_sum_dim_dtype::call(self, dim, dtype);
  11887. }
  11888. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11889. auto results = batch_rule(self_value, self_bdim, dim, dtype);
  11890. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11891. }
  11892. template <typename batch_rule_t, batch_rule_t batch_rule>
  11893. at::Tensor _sparse_sum_backward_generated_plumbing(const at::Tensor & grad, const at::Tensor & self, at::IntArrayRef dim) {
  11894. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11895. auto maybe_layer = maybeCurrentDynamicLayer();
  11896. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11897. int64_t cur_level = maybe_layer->layerId();
  11898. if (!isBatchedAtLevel(grad, cur_level) && !isBatchedAtLevel(self, cur_level)) {
  11899. return at::_ops::_sparse_sum_backward::call(grad, self, dim);
  11900. }
  11901. auto [grad_value, grad_bdim] = unwrapTensorAtLevel(grad, cur_level);
  11902. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11903. auto results = batch_rule(grad_value, grad_bdim, self_value, self_bdim, dim);
  11904. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11905. }
  11906. template <typename batch_rule_t, batch_rule_t batch_rule>
  11907. at::Tensor _sparse_csr_sum_dim_dtype_generated_plumbing(const at::Tensor & self, at::IntArrayRef dim, bool keepdim, ::std::optional<at::ScalarType> dtype) {
  11908. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11909. auto maybe_layer = maybeCurrentDynamicLayer();
  11910. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11911. int64_t cur_level = maybe_layer->layerId();
  11912. if (!isBatchedAtLevel(self, cur_level)) {
  11913. return at::_ops::_sparse_csr_sum_dim_dtype::call(self, dim, keepdim, dtype);
  11914. }
  11915. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11916. auto results = batch_rule(self_value, self_bdim, dim, keepdim, dtype);
  11917. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11918. }
  11919. template <typename batch_rule_t, batch_rule_t batch_rule>
  11920. at::Tensor _sparse_csr_prod_dim_dtype_generated_plumbing(const at::Tensor & self, at::IntArrayRef dim, bool keepdim, ::std::optional<at::ScalarType> dtype) {
  11921. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11922. auto maybe_layer = maybeCurrentDynamicLayer();
  11923. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11924. int64_t cur_level = maybe_layer->layerId();
  11925. if (!isBatchedAtLevel(self, cur_level)) {
  11926. return at::_ops::_sparse_csr_prod_dim_dtype::call(self, dim, keepdim, dtype);
  11927. }
  11928. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11929. auto results = batch_rule(self_value, self_bdim, dim, keepdim, dtype);
  11930. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11931. }
  11932. template <typename batch_rule_t, batch_rule_t batch_rule>
  11933. at::Tensor _sparse_softmax_int_generated_plumbing(const at::Tensor & self, int64_t dim, ::std::optional<at::ScalarType> dtype) {
  11934. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11935. auto maybe_layer = maybeCurrentDynamicLayer();
  11936. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11937. int64_t cur_level = maybe_layer->layerId();
  11938. if (!isBatchedAtLevel(self, cur_level)) {
  11939. return at::_ops::_sparse_softmax_int::call(self, dim, dtype);
  11940. }
  11941. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11942. auto results = batch_rule(self_value, self_bdim, dim, dtype);
  11943. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11944. }
  11945. template <typename batch_rule_t, batch_rule_t batch_rule>
  11946. at::Tensor _sparse_softmax_Dimname_generated_plumbing(const at::Tensor & self, at::Dimname dim, ::std::optional<at::ScalarType> dtype) {
  11947. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11948. auto maybe_layer = maybeCurrentDynamicLayer();
  11949. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11950. int64_t cur_level = maybe_layer->layerId();
  11951. if (!isBatchedAtLevel(self, cur_level)) {
  11952. return at::_ops::_sparse_softmax_Dimname::call(self, dim, dtype);
  11953. }
  11954. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11955. auto results = batch_rule(self_value, self_bdim, dim, dtype);
  11956. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11957. }
  11958. template <typename batch_rule_t, batch_rule_t batch_rule>
  11959. at::Tensor _sparse_softmax_generated_plumbing(const at::Tensor & self, int64_t dim, bool half_to_float) {
  11960. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11961. auto maybe_layer = maybeCurrentDynamicLayer();
  11962. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11963. int64_t cur_level = maybe_layer->layerId();
  11964. if (!isBatchedAtLevel(self, cur_level)) {
  11965. return at::_ops::_sparse_softmax::call(self, dim, half_to_float);
  11966. }
  11967. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11968. auto results = batch_rule(self_value, self_bdim, dim, half_to_float);
  11969. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11970. }
  11971. template <typename batch_rule_t, batch_rule_t batch_rule>
  11972. at::Tensor _sparse_softmax_backward_data_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & output, int64_t dim, const at::Tensor & self) {
  11973. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11974. auto maybe_layer = maybeCurrentDynamicLayer();
  11975. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11976. int64_t cur_level = maybe_layer->layerId();
  11977. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(output, cur_level) && !isBatchedAtLevel(self, cur_level)) {
  11978. return at::_ops::_sparse_softmax_backward_data::call(grad_output, output, dim, self);
  11979. }
  11980. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  11981. auto [output_value, output_bdim] = unwrapTensorAtLevel(output, cur_level);
  11982. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11983. auto results = batch_rule(grad_output_value, grad_output_bdim, output_value, output_bdim, dim, self_value, self_bdim);
  11984. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11985. }
  11986. template <typename batch_rule_t, batch_rule_t batch_rule>
  11987. at::Tensor _sparse_log_softmax_int_generated_plumbing(const at::Tensor & self, int64_t dim, ::std::optional<at::ScalarType> dtype) {
  11988. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  11989. auto maybe_layer = maybeCurrentDynamicLayer();
  11990. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  11991. int64_t cur_level = maybe_layer->layerId();
  11992. if (!isBatchedAtLevel(self, cur_level)) {
  11993. return at::_ops::_sparse_log_softmax_int::call(self, dim, dtype);
  11994. }
  11995. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  11996. auto results = batch_rule(self_value, self_bdim, dim, dtype);
  11997. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  11998. }
  11999. template <typename batch_rule_t, batch_rule_t batch_rule>
  12000. at::Tensor _sparse_log_softmax_Dimname_generated_plumbing(const at::Tensor & self, at::Dimname dim, ::std::optional<at::ScalarType> dtype) {
  12001. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12002. auto maybe_layer = maybeCurrentDynamicLayer();
  12003. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12004. int64_t cur_level = maybe_layer->layerId();
  12005. if (!isBatchedAtLevel(self, cur_level)) {
  12006. return at::_ops::_sparse_log_softmax_Dimname::call(self, dim, dtype);
  12007. }
  12008. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  12009. auto results = batch_rule(self_value, self_bdim, dim, dtype);
  12010. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12011. }
  12012. template <typename batch_rule_t, batch_rule_t batch_rule>
  12013. at::Tensor _sparse_log_softmax_generated_plumbing(const at::Tensor & self, int64_t dim, bool half_to_float) {
  12014. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12015. auto maybe_layer = maybeCurrentDynamicLayer();
  12016. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12017. int64_t cur_level = maybe_layer->layerId();
  12018. if (!isBatchedAtLevel(self, cur_level)) {
  12019. return at::_ops::_sparse_log_softmax::call(self, dim, half_to_float);
  12020. }
  12021. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  12022. auto results = batch_rule(self_value, self_bdim, dim, half_to_float);
  12023. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12024. }
  12025. template <typename batch_rule_t, batch_rule_t batch_rule>
  12026. at::Tensor _sparse_log_softmax_backward_data_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & output, int64_t dim, const at::Tensor & self) {
  12027. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12028. auto maybe_layer = maybeCurrentDynamicLayer();
  12029. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12030. int64_t cur_level = maybe_layer->layerId();
  12031. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(output, cur_level) && !isBatchedAtLevel(self, cur_level)) {
  12032. return at::_ops::_sparse_log_softmax_backward_data::call(grad_output, output, dim, self);
  12033. }
  12034. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  12035. auto [output_value, output_bdim] = unwrapTensorAtLevel(output, cur_level);
  12036. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  12037. auto results = batch_rule(grad_output_value, grad_output_bdim, output_value, output_bdim, dim, self_value, self_bdim);
  12038. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12039. }
  12040. template <typename batch_rule_t, batch_rule_t batch_rule>
  12041. at::Tensor _spdiags_generated_plumbing(const at::Tensor & diagonals, const at::Tensor & offsets, at::IntArrayRef shape, ::std::optional<at::Layout> layout) {
  12042. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12043. auto maybe_layer = maybeCurrentDynamicLayer();
  12044. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12045. int64_t cur_level = maybe_layer->layerId();
  12046. if (!isBatchedAtLevel(diagonals, cur_level) && !isBatchedAtLevel(offsets, cur_level)) {
  12047. return at::_ops::_spdiags::call(diagonals, offsets, shape, layout);
  12048. }
  12049. auto [diagonals_value, diagonals_bdim] = unwrapTensorAtLevel(diagonals, cur_level);
  12050. auto [offsets_value, offsets_bdim] = unwrapTensorAtLevel(offsets, cur_level);
  12051. auto results = batch_rule(diagonals_value, diagonals_bdim, offsets_value, offsets_bdim, shape, layout);
  12052. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12053. }
  12054. template <typename batch_rule_t, batch_rule_t batch_rule>
  12055. at::Tensor norm_ScalarOpt_dtype_generated_plumbing(const at::Tensor & self, const ::std::optional<at::Scalar> & p, at::ScalarType dtype) {
  12056. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12057. auto maybe_layer = maybeCurrentDynamicLayer();
  12058. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12059. int64_t cur_level = maybe_layer->layerId();
  12060. if (!isBatchedAtLevel(self, cur_level)) {
  12061. return at::_ops::norm_ScalarOpt_dtype::call(self, p, dtype);
  12062. }
  12063. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  12064. auto results = batch_rule(self_value, self_bdim, p, dtype);
  12065. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12066. }
  12067. template <typename batch_rule_t, batch_rule_t batch_rule>
  12068. at::Tensor norm_Scalar_generated_plumbing(const at::Tensor & self, const at::Scalar & p) {
  12069. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12070. auto maybe_layer = maybeCurrentDynamicLayer();
  12071. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12072. int64_t cur_level = maybe_layer->layerId();
  12073. if (!isBatchedAtLevel(self, cur_level)) {
  12074. return at::_ops::norm_Scalar::call(self, p);
  12075. }
  12076. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  12077. auto results = batch_rule(self_value, self_bdim, p);
  12078. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12079. }
  12080. template <typename batch_rule_t, batch_rule_t batch_rule>
  12081. at::Tensor norm_ScalarOpt_dim_dtype_generated_plumbing(const at::Tensor & self, const ::std::optional<at::Scalar> & p, at::IntArrayRef dim, bool keepdim, at::ScalarType dtype) {
  12082. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12083. auto maybe_layer = maybeCurrentDynamicLayer();
  12084. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12085. int64_t cur_level = maybe_layer->layerId();
  12086. if (!isBatchedAtLevel(self, cur_level)) {
  12087. return at::_ops::norm_ScalarOpt_dim_dtype::call(self, p, dim, keepdim, dtype);
  12088. }
  12089. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  12090. auto results = batch_rule(self_value, self_bdim, p, dim, keepdim, dtype);
  12091. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12092. }
  12093. template <typename batch_rule_t, batch_rule_t batch_rule>
  12094. at::Tensor norm_ScalarOpt_dim_generated_plumbing(const at::Tensor & self, const ::std::optional<at::Scalar> & p, at::IntArrayRef dim, bool keepdim) {
  12095. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12096. auto maybe_layer = maybeCurrentDynamicLayer();
  12097. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12098. int64_t cur_level = maybe_layer->layerId();
  12099. if (!isBatchedAtLevel(self, cur_level)) {
  12100. return at::_ops::norm_ScalarOpt_dim::call(self, p, dim, keepdim);
  12101. }
  12102. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  12103. auto results = batch_rule(self_value, self_bdim, p, dim, keepdim);
  12104. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12105. }
  12106. template <typename batch_rule_t, batch_rule_t batch_rule>
  12107. at::Tensor norm_names_ScalarOpt_dim_dtype_generated_plumbing(const at::Tensor & self, const ::std::optional<at::Scalar> & p, at::DimnameList dim, bool keepdim, at::ScalarType dtype) {
  12108. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12109. auto maybe_layer = maybeCurrentDynamicLayer();
  12110. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12111. int64_t cur_level = maybe_layer->layerId();
  12112. if (!isBatchedAtLevel(self, cur_level)) {
  12113. return at::_ops::norm_names_ScalarOpt_dim_dtype::call(self, p, dim, keepdim, dtype);
  12114. }
  12115. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  12116. auto results = batch_rule(self_value, self_bdim, p, dim, keepdim, dtype);
  12117. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12118. }
  12119. template <typename batch_rule_t, batch_rule_t batch_rule>
  12120. at::Tensor norm_names_ScalarOpt_dim_generated_plumbing(const at::Tensor & self, const ::std::optional<at::Scalar> & p, at::DimnameList dim, bool keepdim) {
  12121. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12122. auto maybe_layer = maybeCurrentDynamicLayer();
  12123. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12124. int64_t cur_level = maybe_layer->layerId();
  12125. if (!isBatchedAtLevel(self, cur_level)) {
  12126. return at::_ops::norm_names_ScalarOpt_dim::call(self, p, dim, keepdim);
  12127. }
  12128. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  12129. auto results = batch_rule(self_value, self_bdim, p, dim, keepdim);
  12130. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12131. }
  12132. template <typename batch_rule_t, batch_rule_t batch_rule>
  12133. ::std::tuple<at::Tensor,at::Tensor> frexp_Tensor_generated_plumbing(const at::Tensor & self) {
  12134. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12135. auto maybe_layer = maybeCurrentDynamicLayer();
  12136. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12137. int64_t cur_level = maybe_layer->layerId();
  12138. if (!isBatchedAtLevel(self, cur_level)) {
  12139. return at::_ops::frexp_Tensor::call(self);
  12140. }
  12141. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  12142. auto results = batch_rule(self_value, self_bdim);
  12143. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  12144. }
  12145. template <typename batch_rule_t, batch_rule_t batch_rule>
  12146. at::Tensor frobenius_norm_dim_generated_plumbing(const at::Tensor & self, at::IntArrayRef dim, bool keepdim) {
  12147. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12148. auto maybe_layer = maybeCurrentDynamicLayer();
  12149. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12150. int64_t cur_level = maybe_layer->layerId();
  12151. if (!isBatchedAtLevel(self, cur_level)) {
  12152. return at::_ops::frobenius_norm_dim::call(self, dim, keepdim);
  12153. }
  12154. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  12155. auto results = batch_rule(self_value, self_bdim, dim, keepdim);
  12156. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12157. }
  12158. template <typename batch_rule_t, batch_rule_t batch_rule>
  12159. at::Tensor nuclear_norm_generated_plumbing(const at::Tensor & self, bool keepdim) {
  12160. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12161. auto maybe_layer = maybeCurrentDynamicLayer();
  12162. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12163. int64_t cur_level = maybe_layer->layerId();
  12164. if (!isBatchedAtLevel(self, cur_level)) {
  12165. return at::_ops::nuclear_norm::call(self, keepdim);
  12166. }
  12167. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  12168. auto results = batch_rule(self_value, self_bdim, keepdim);
  12169. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12170. }
  12171. template <typename batch_rule_t, batch_rule_t batch_rule>
  12172. at::Tensor nuclear_norm_dim_generated_plumbing(const at::Tensor & self, at::IntArrayRef dim, bool keepdim) {
  12173. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12174. auto maybe_layer = maybeCurrentDynamicLayer();
  12175. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12176. int64_t cur_level = maybe_layer->layerId();
  12177. if (!isBatchedAtLevel(self, cur_level)) {
  12178. return at::_ops::nuclear_norm_dim::call(self, dim, keepdim);
  12179. }
  12180. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  12181. auto results = batch_rule(self_value, self_bdim, dim, keepdim);
  12182. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12183. }
  12184. template <typename batch_rule_t, batch_rule_t batch_rule>
  12185. at::Tensor clone_generated_plumbing(const at::Tensor & self, ::std::optional<at::MemoryFormat> memory_format) {
  12186. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12187. auto maybe_layer = maybeCurrentDynamicLayer();
  12188. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12189. int64_t cur_level = maybe_layer->layerId();
  12190. if (!isBatchedAtLevel(self, cur_level)) {
  12191. return at::_ops::clone::call(self, memory_format);
  12192. }
  12193. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  12194. auto results = batch_rule(self_value, self_bdim, memory_format);
  12195. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12196. }
  12197. template <typename batch_rule_t, batch_rule_t batch_rule>
  12198. at::Tensor positive_generated_plumbing(const at::Tensor & self) {
  12199. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12200. auto maybe_layer = maybeCurrentDynamicLayer();
  12201. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12202. int64_t cur_level = maybe_layer->layerId();
  12203. if (!isBatchedAtLevel(self, cur_level)) {
  12204. return at::_ops::positive::call(self);
  12205. }
  12206. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  12207. auto results = batch_rule(self_value, self_bdim);
  12208. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12209. }
  12210. template <typename batch_rule_t, batch_rule_t batch_rule>
  12211. const at::Tensor & resize_as_sparse__generated_plumbing(const at::Tensor & self, const at::Tensor & the_template) {
  12212. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12213. auto maybe_layer = maybeCurrentDynamicLayer();
  12214. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  12215. int64_t cur_level = maybe_layer->layerId();
  12216. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(the_template, cur_level)) {
  12217. return at::_ops::resize_as_sparse_::call(self, the_template);
  12218. }
  12219. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  12220. auto [the_template_value, the_template_bdim] = unwrapTensorAtLevel(the_template, cur_level);
  12221. batch_rule(self_value, self_bdim, the_template_value, the_template_bdim);
  12222. return self;
  12223. }
  12224. template <typename batch_rule_t, batch_rule_t batch_rule>
  12225. at::Tensor & zero__generated_plumbing(at::Tensor & self) {
  12226. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12227. auto maybe_layer = maybeCurrentDynamicLayer();
  12228. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  12229. int64_t cur_level = maybe_layer->layerId();
  12230. if (!isBatchedAtLevel(self, cur_level)) {
  12231. return at::_ops::zero_::call(self);
  12232. }
  12233. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  12234. batch_rule(self_value, self_bdim);
  12235. return self;
  12236. }
  12237. template <typename batch_rule_t, batch_rule_t batch_rule>
  12238. at::Tensor sub_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & other, const at::Scalar & alpha) {
  12239. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12240. auto maybe_layer = maybeCurrentDynamicLayer();
  12241. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12242. int64_t cur_level = maybe_layer->layerId();
  12243. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  12244. return at::_ops::sub_Tensor::call(self, other, alpha);
  12245. }
  12246. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  12247. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  12248. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim, alpha);
  12249. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12250. }
  12251. template <typename batch_rule_t, batch_rule_t batch_rule>
  12252. at::Tensor & sub__Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & other, const at::Scalar & alpha) {
  12253. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12254. auto maybe_layer = maybeCurrentDynamicLayer();
  12255. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  12256. int64_t cur_level = maybe_layer->layerId();
  12257. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  12258. return at::_ops::sub__Tensor::call(self, other, alpha);
  12259. }
  12260. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  12261. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  12262. batch_rule(self_value, self_bdim, other_value, other_bdim, alpha);
  12263. return self;
  12264. }
  12265. template <typename batch_rule_t, batch_rule_t batch_rule>
  12266. at::Tensor sub_Scalar_generated_plumbing(const at::Tensor & self, const at::Scalar & other, const at::Scalar & alpha) {
  12267. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12268. auto maybe_layer = maybeCurrentDynamicLayer();
  12269. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12270. int64_t cur_level = maybe_layer->layerId();
  12271. if (!isBatchedAtLevel(self, cur_level)) {
  12272. return at::_ops::sub_Scalar::call(self, other, alpha);
  12273. }
  12274. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  12275. auto results = batch_rule(self_value, self_bdim, other, alpha);
  12276. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12277. }
  12278. template <typename batch_rule_t, batch_rule_t batch_rule>
  12279. at::Tensor & sub__Scalar_generated_plumbing(at::Tensor & self, const at::Scalar & other, const at::Scalar & alpha) {
  12280. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12281. auto maybe_layer = maybeCurrentDynamicLayer();
  12282. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  12283. int64_t cur_level = maybe_layer->layerId();
  12284. if (!isBatchedAtLevel(self, cur_level)) {
  12285. return at::_ops::sub__Scalar::call(self, other, alpha);
  12286. }
  12287. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  12288. batch_rule(self_value, self_bdim, other, alpha);
  12289. return self;
  12290. }
  12291. template <typename batch_rule_t, batch_rule_t batch_rule>
  12292. at::Tensor subtract_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & other, const at::Scalar & alpha) {
  12293. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12294. auto maybe_layer = maybeCurrentDynamicLayer();
  12295. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12296. int64_t cur_level = maybe_layer->layerId();
  12297. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  12298. return at::_ops::subtract_Tensor::call(self, other, alpha);
  12299. }
  12300. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  12301. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  12302. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim, alpha);
  12303. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12304. }
  12305. template <typename batch_rule_t, batch_rule_t batch_rule>
  12306. at::Tensor & subtract__Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & other, const at::Scalar & alpha) {
  12307. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12308. auto maybe_layer = maybeCurrentDynamicLayer();
  12309. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  12310. int64_t cur_level = maybe_layer->layerId();
  12311. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  12312. return at::_ops::subtract__Tensor::call(self, other, alpha);
  12313. }
  12314. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  12315. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  12316. batch_rule(self_value, self_bdim, other_value, other_bdim, alpha);
  12317. return self;
  12318. }
  12319. template <typename batch_rule_t, batch_rule_t batch_rule>
  12320. at::Tensor subtract_Scalar_generated_plumbing(const at::Tensor & self, const at::Scalar & other, const at::Scalar & alpha) {
  12321. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12322. auto maybe_layer = maybeCurrentDynamicLayer();
  12323. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12324. int64_t cur_level = maybe_layer->layerId();
  12325. if (!isBatchedAtLevel(self, cur_level)) {
  12326. return at::_ops::subtract_Scalar::call(self, other, alpha);
  12327. }
  12328. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  12329. auto results = batch_rule(self_value, self_bdim, other, alpha);
  12330. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12331. }
  12332. template <typename batch_rule_t, batch_rule_t batch_rule>
  12333. at::Tensor & subtract__Scalar_generated_plumbing(at::Tensor & self, const at::Scalar & other, const at::Scalar & alpha) {
  12334. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12335. auto maybe_layer = maybeCurrentDynamicLayer();
  12336. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  12337. int64_t cur_level = maybe_layer->layerId();
  12338. if (!isBatchedAtLevel(self, cur_level)) {
  12339. return at::_ops::subtract__Scalar::call(self, other, alpha);
  12340. }
  12341. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  12342. batch_rule(self_value, self_bdim, other, alpha);
  12343. return self;
  12344. }
  12345. template <typename batch_rule_t, batch_rule_t batch_rule>
  12346. at::Tensor rsub_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & other, const at::Scalar & alpha) {
  12347. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12348. auto maybe_layer = maybeCurrentDynamicLayer();
  12349. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12350. int64_t cur_level = maybe_layer->layerId();
  12351. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  12352. return at::_ops::rsub_Tensor::call(self, other, alpha);
  12353. }
  12354. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  12355. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  12356. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim, alpha);
  12357. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12358. }
  12359. template <typename batch_rule_t, batch_rule_t batch_rule>
  12360. at::Tensor heaviside_generated_plumbing(const at::Tensor & self, const at::Tensor & values) {
  12361. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12362. auto maybe_layer = maybeCurrentDynamicLayer();
  12363. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12364. int64_t cur_level = maybe_layer->layerId();
  12365. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(values, cur_level)) {
  12366. return at::_ops::heaviside::call(self, values);
  12367. }
  12368. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  12369. auto [values_value, values_bdim] = unwrapTensorAtLevel(values, cur_level);
  12370. auto results = batch_rule(self_value, self_bdim, values_value, values_bdim);
  12371. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12372. }
  12373. template <typename batch_rule_t, batch_rule_t batch_rule>
  12374. at::Tensor & heaviside__generated_plumbing(at::Tensor & self, const at::Tensor & values) {
  12375. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12376. auto maybe_layer = maybeCurrentDynamicLayer();
  12377. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  12378. int64_t cur_level = maybe_layer->layerId();
  12379. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(values, cur_level)) {
  12380. return at::_ops::heaviside_::call(self, values);
  12381. }
  12382. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  12383. auto [values_value, values_bdim] = unwrapTensorAtLevel(values, cur_level);
  12384. batch_rule(self_value, self_bdim, values_value, values_bdim);
  12385. return self;
  12386. }
  12387. template <typename batch_rule_t, batch_rule_t batch_rule>
  12388. at::Tensor rsub_Scalar_generated_plumbing(const at::Tensor & self, const at::Scalar & other, const at::Scalar & alpha) {
  12389. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12390. auto maybe_layer = maybeCurrentDynamicLayer();
  12391. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12392. int64_t cur_level = maybe_layer->layerId();
  12393. if (!isBatchedAtLevel(self, cur_level)) {
  12394. return at::_ops::rsub_Scalar::call(self, other, alpha);
  12395. }
  12396. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  12397. auto results = batch_rule(self_value, self_bdim, other, alpha);
  12398. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12399. }
  12400. template <typename batch_rule_t, batch_rule_t batch_rule>
  12401. at::Tensor _sparse_addmm_generated_plumbing(const at::Tensor & self, const at::Tensor & mat1, const at::Tensor & mat2, const at::Scalar & beta, const at::Scalar & alpha) {
  12402. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12403. auto maybe_layer = maybeCurrentDynamicLayer();
  12404. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12405. int64_t cur_level = maybe_layer->layerId();
  12406. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(mat1, cur_level) && !isBatchedAtLevel(mat2, cur_level)) {
  12407. return at::_ops::_sparse_addmm::call(self, mat1, mat2, beta, alpha);
  12408. }
  12409. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  12410. auto [mat1_value, mat1_bdim] = unwrapTensorAtLevel(mat1, cur_level);
  12411. auto [mat2_value, mat2_bdim] = unwrapTensorAtLevel(mat2, cur_level);
  12412. auto results = batch_rule(self_value, self_bdim, mat1_value, mat1_bdim, mat2_value, mat2_bdim, beta, alpha);
  12413. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12414. }
  12415. template <typename batch_rule_t, batch_rule_t batch_rule>
  12416. at::Tensor sparse_sampled_addmm_generated_plumbing(const at::Tensor & self, const at::Tensor & mat1, const at::Tensor & mat2, const at::Scalar & beta, const at::Scalar & alpha) {
  12417. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12418. auto maybe_layer = maybeCurrentDynamicLayer();
  12419. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12420. int64_t cur_level = maybe_layer->layerId();
  12421. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(mat1, cur_level) && !isBatchedAtLevel(mat2, cur_level)) {
  12422. return at::_ops::sparse_sampled_addmm::call(self, mat1, mat2, beta, alpha);
  12423. }
  12424. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  12425. auto [mat1_value, mat1_bdim] = unwrapTensorAtLevel(mat1, cur_level);
  12426. auto [mat2_value, mat2_bdim] = unwrapTensorAtLevel(mat2, cur_level);
  12427. auto results = batch_rule(self_value, self_bdim, mat1_value, mat1_bdim, mat2_value, mat2_bdim, beta, alpha);
  12428. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12429. }
  12430. template <typename batch_rule_t, batch_rule_t batch_rule>
  12431. ::std::tuple<at::Tensor,at::Tensor> _sparse_mm_reduce_impl_generated_plumbing(const at::Tensor & self, const at::Tensor & other, c10::string_view reduce) {
  12432. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12433. auto maybe_layer = maybeCurrentDynamicLayer();
  12434. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12435. int64_t cur_level = maybe_layer->layerId();
  12436. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  12437. return at::_ops::_sparse_mm_reduce_impl::call(self, other, reduce);
  12438. }
  12439. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  12440. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  12441. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim, reduce);
  12442. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  12443. }
  12444. template <typename batch_rule_t, batch_rule_t batch_rule>
  12445. ::std::tuple<at::Tensor,at::Tensor> _sparse_mm_reduce_impl_backward_generated_plumbing(const at::Tensor & self, const at::Tensor & grad_out, const at::Tensor & weight, c10::string_view reduce, const at::Tensor & arg_out, ::std::array<bool,2> output_mask) {
  12446. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12447. auto maybe_layer = maybeCurrentDynamicLayer();
  12448. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12449. int64_t cur_level = maybe_layer->layerId();
  12450. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(grad_out, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(arg_out, cur_level)) {
  12451. return at::_ops::_sparse_mm_reduce_impl_backward::call(self, grad_out, weight, reduce, arg_out, output_mask);
  12452. }
  12453. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  12454. auto [grad_out_value, grad_out_bdim] = unwrapTensorAtLevel(grad_out, cur_level);
  12455. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  12456. auto [arg_out_value, arg_out_bdim] = unwrapTensorAtLevel(arg_out, cur_level);
  12457. auto results = batch_rule(self_value, self_bdim, grad_out_value, grad_out_bdim, weight_value, weight_bdim, reduce, arg_out_value, arg_out_bdim, output_mask);
  12458. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  12459. }
  12460. template <typename batch_rule_t, batch_rule_t batch_rule>
  12461. at::Tensor addmm_generated_plumbing(const at::Tensor & self, const at::Tensor & mat1, const at::Tensor & mat2, const at::Scalar & beta, const at::Scalar & alpha) {
  12462. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12463. auto maybe_layer = maybeCurrentDynamicLayer();
  12464. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12465. int64_t cur_level = maybe_layer->layerId();
  12466. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(mat1, cur_level) && !isBatchedAtLevel(mat2, cur_level)) {
  12467. return at::_ops::addmm::call(self, mat1, mat2, beta, alpha);
  12468. }
  12469. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  12470. auto [mat1_value, mat1_bdim] = unwrapTensorAtLevel(mat1, cur_level);
  12471. auto [mat2_value, mat2_bdim] = unwrapTensorAtLevel(mat2, cur_level);
  12472. auto results = batch_rule(self_value, self_bdim, mat1_value, mat1_bdim, mat2_value, mat2_bdim, beta, alpha);
  12473. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12474. }
  12475. template <typename batch_rule_t, batch_rule_t batch_rule>
  12476. at::Tensor addmm_dtype_generated_plumbing(const at::Tensor & self, const at::Tensor & mat1, const at::Tensor & mat2, at::ScalarType out_dtype, const at::Scalar & beta, const at::Scalar & alpha) {
  12477. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12478. auto maybe_layer = maybeCurrentDynamicLayer();
  12479. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12480. int64_t cur_level = maybe_layer->layerId();
  12481. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(mat1, cur_level) && !isBatchedAtLevel(mat2, cur_level)) {
  12482. return at::_ops::addmm_dtype::call(self, mat1, mat2, out_dtype, beta, alpha);
  12483. }
  12484. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  12485. auto [mat1_value, mat1_bdim] = unwrapTensorAtLevel(mat1, cur_level);
  12486. auto [mat2_value, mat2_bdim] = unwrapTensorAtLevel(mat2, cur_level);
  12487. auto results = batch_rule(self_value, self_bdim, mat1_value, mat1_bdim, mat2_value, mat2_bdim, out_dtype, beta, alpha);
  12488. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12489. }
  12490. template <typename batch_rule_t, batch_rule_t batch_rule>
  12491. at::Tensor & addmm__generated_plumbing(at::Tensor & self, const at::Tensor & mat1, const at::Tensor & mat2, const at::Scalar & beta, const at::Scalar & alpha) {
  12492. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12493. auto maybe_layer = maybeCurrentDynamicLayer();
  12494. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  12495. int64_t cur_level = maybe_layer->layerId();
  12496. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(mat1, cur_level) && !isBatchedAtLevel(mat2, cur_level)) {
  12497. return at::_ops::addmm_::call(self, mat1, mat2, beta, alpha);
  12498. }
  12499. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  12500. auto [mat1_value, mat1_bdim] = unwrapTensorAtLevel(mat1, cur_level);
  12501. auto [mat2_value, mat2_bdim] = unwrapTensorAtLevel(mat2, cur_level);
  12502. batch_rule(self_value, self_bdim, mat1_value, mat1_bdim, mat2_value, mat2_bdim, beta, alpha);
  12503. return self;
  12504. }
  12505. template <typename batch_rule_t, batch_rule_t batch_rule>
  12506. at::Tensor _addmm_activation_generated_plumbing(const at::Tensor & self, const at::Tensor & mat1, const at::Tensor & mat2, const at::Scalar & beta, const at::Scalar & alpha, bool use_gelu) {
  12507. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12508. auto maybe_layer = maybeCurrentDynamicLayer();
  12509. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12510. int64_t cur_level = maybe_layer->layerId();
  12511. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(mat1, cur_level) && !isBatchedAtLevel(mat2, cur_level)) {
  12512. return at::_ops::_addmm_activation::call(self, mat1, mat2, beta, alpha, use_gelu);
  12513. }
  12514. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  12515. auto [mat1_value, mat1_bdim] = unwrapTensorAtLevel(mat1, cur_level);
  12516. auto [mat2_value, mat2_bdim] = unwrapTensorAtLevel(mat2, cur_level);
  12517. auto results = batch_rule(self_value, self_bdim, mat1_value, mat1_bdim, mat2_value, mat2_bdim, beta, alpha, use_gelu);
  12518. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12519. }
  12520. template <typename batch_rule_t, batch_rule_t batch_rule>
  12521. at::Tensor _scaled_mm_generated_plumbing(const at::Tensor & self, const at::Tensor & mat2, const at::Tensor & scale_a, const at::Tensor & scale_b, const ::std::optional<at::Tensor> & bias, const ::std::optional<at::Tensor> & scale_result, ::std::optional<at::ScalarType> out_dtype, bool use_fast_accum) {
  12522. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12523. auto maybe_layer = maybeCurrentDynamicLayer();
  12524. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12525. int64_t cur_level = maybe_layer->layerId();
  12526. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(mat2, cur_level) && !isBatchedAtLevel(scale_a, cur_level) && !isBatchedAtLevel(scale_b, cur_level) && !isBatchedAtLevel(bias, cur_level) && !isBatchedAtLevel(scale_result, cur_level)) {
  12527. return at::_ops::_scaled_mm::call(self, mat2, scale_a, scale_b, bias, scale_result, out_dtype, use_fast_accum);
  12528. }
  12529. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  12530. auto [mat2_value, mat2_bdim] = unwrapTensorAtLevel(mat2, cur_level);
  12531. auto [scale_a_value, scale_a_bdim] = unwrapTensorAtLevel(scale_a, cur_level);
  12532. auto [scale_b_value, scale_b_bdim] = unwrapTensorAtLevel(scale_b, cur_level);
  12533. std::optional<Tensor> bias_value;
  12534. std::optional<int64_t> bias_bdim;
  12535. if (bias) {
  12536. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  12537. }
  12538. std::optional<Tensor> scale_result_value;
  12539. std::optional<int64_t> scale_result_bdim;
  12540. if (scale_result) {
  12541. std::tie(scale_result_value, scale_result_bdim) = unwrapTensorAtLevel(scale_result.value(), cur_level);
  12542. }
  12543. auto results = batch_rule(self_value, self_bdim, mat2_value, mat2_bdim, scale_a_value, scale_a_bdim, scale_b_value, scale_b_bdim, bias_value, bias_bdim, scale_result_value, scale_result_bdim, out_dtype, use_fast_accum);
  12544. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12545. }
  12546. template <typename batch_rule_t, batch_rule_t batch_rule>
  12547. at::Tensor _scaled_grouped_mm_generated_plumbing(const at::Tensor & self, const at::Tensor & mat2, const at::Tensor & scale_a, const at::Tensor & scale_b, const ::std::optional<at::Tensor> & offs, const ::std::optional<at::Tensor> & bias, const ::std::optional<at::Tensor> & scale_result, ::std::optional<at::ScalarType> out_dtype, bool use_fast_accum) {
  12548. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12549. auto maybe_layer = maybeCurrentDynamicLayer();
  12550. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12551. int64_t cur_level = maybe_layer->layerId();
  12552. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(mat2, cur_level) && !isBatchedAtLevel(scale_a, cur_level) && !isBatchedAtLevel(scale_b, cur_level) && !isBatchedAtLevel(offs, cur_level) && !isBatchedAtLevel(bias, cur_level) && !isBatchedAtLevel(scale_result, cur_level)) {
  12553. return at::_ops::_scaled_grouped_mm::call(self, mat2, scale_a, scale_b, offs, bias, scale_result, out_dtype, use_fast_accum);
  12554. }
  12555. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  12556. auto [mat2_value, mat2_bdim] = unwrapTensorAtLevel(mat2, cur_level);
  12557. auto [scale_a_value, scale_a_bdim] = unwrapTensorAtLevel(scale_a, cur_level);
  12558. auto [scale_b_value, scale_b_bdim] = unwrapTensorAtLevel(scale_b, cur_level);
  12559. std::optional<Tensor> offs_value;
  12560. std::optional<int64_t> offs_bdim;
  12561. if (offs) {
  12562. std::tie(offs_value, offs_bdim) = unwrapTensorAtLevel(offs.value(), cur_level);
  12563. }
  12564. std::optional<Tensor> bias_value;
  12565. std::optional<int64_t> bias_bdim;
  12566. if (bias) {
  12567. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  12568. }
  12569. std::optional<Tensor> scale_result_value;
  12570. std::optional<int64_t> scale_result_bdim;
  12571. if (scale_result) {
  12572. std::tie(scale_result_value, scale_result_bdim) = unwrapTensorAtLevel(scale_result.value(), cur_level);
  12573. }
  12574. auto results = batch_rule(self_value, self_bdim, mat2_value, mat2_bdim, scale_a_value, scale_a_bdim, scale_b_value, scale_b_bdim, offs_value, offs_bdim, bias_value, bias_bdim, scale_result_value, scale_result_bdim, out_dtype, use_fast_accum);
  12575. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12576. }
  12577. template <typename batch_rule_t, batch_rule_t batch_rule>
  12578. at::Tensor _grouped_mm_generated_plumbing(const at::Tensor & self, const at::Tensor & mat2, const ::std::optional<at::Tensor> & offs, const ::std::optional<at::Tensor> & bias, ::std::optional<at::ScalarType> out_dtype) {
  12579. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12580. auto maybe_layer = maybeCurrentDynamicLayer();
  12581. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12582. int64_t cur_level = maybe_layer->layerId();
  12583. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(mat2, cur_level) && !isBatchedAtLevel(offs, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  12584. return at::_ops::_grouped_mm::call(self, mat2, offs, bias, out_dtype);
  12585. }
  12586. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  12587. auto [mat2_value, mat2_bdim] = unwrapTensorAtLevel(mat2, cur_level);
  12588. std::optional<Tensor> offs_value;
  12589. std::optional<int64_t> offs_bdim;
  12590. if (offs) {
  12591. std::tie(offs_value, offs_bdim) = unwrapTensorAtLevel(offs.value(), cur_level);
  12592. }
  12593. std::optional<Tensor> bias_value;
  12594. std::optional<int64_t> bias_bdim;
  12595. if (bias) {
  12596. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  12597. }
  12598. auto results = batch_rule(self_value, self_bdim, mat2_value, mat2_bdim, offs_value, offs_bdim, bias_value, bias_bdim, out_dtype);
  12599. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12600. }
  12601. template <typename batch_rule_t, batch_rule_t batch_rule>
  12602. at::Tensor sparse_compressed_tensor_comp_plain_value_size_generated_plumbing(const at::Tensor & compressed_indices, const at::Tensor & plain_indices, const at::Tensor & values, c10::SymIntArrayRef size, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory) {
  12603. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12604. auto maybe_layer = maybeCurrentDynamicLayer();
  12605. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12606. int64_t cur_level = maybe_layer->layerId();
  12607. if (!isBatchedAtLevel(compressed_indices, cur_level) && !isBatchedAtLevel(plain_indices, cur_level) && !isBatchedAtLevel(values, cur_level)) {
  12608. return at::_ops::sparse_compressed_tensor_comp_plain_value_size::call(compressed_indices, plain_indices, values, size, dtype, layout, device, pin_memory);
  12609. }
  12610. auto [compressed_indices_value, compressed_indices_bdim] = unwrapTensorAtLevel(compressed_indices, cur_level);
  12611. auto [plain_indices_value, plain_indices_bdim] = unwrapTensorAtLevel(plain_indices, cur_level);
  12612. auto [values_value, values_bdim] = unwrapTensorAtLevel(values, cur_level);
  12613. auto results = batch_rule(compressed_indices_value, compressed_indices_bdim, plain_indices_value, plain_indices_bdim, values_value, values_bdim, size, dtype, layout, device, pin_memory);
  12614. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12615. }
  12616. template <typename batch_rule_t, batch_rule_t batch_rule>
  12617. at::Tensor sparse_csr_tensor_crow_col_value_size_generated_plumbing(const at::Tensor & crow_indices, const at::Tensor & col_indices, const at::Tensor & values, at::IntArrayRef size, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory) {
  12618. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12619. auto maybe_layer = maybeCurrentDynamicLayer();
  12620. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12621. int64_t cur_level = maybe_layer->layerId();
  12622. if (!isBatchedAtLevel(crow_indices, cur_level) && !isBatchedAtLevel(col_indices, cur_level) && !isBatchedAtLevel(values, cur_level)) {
  12623. return at::_ops::sparse_csr_tensor_crow_col_value_size::call(crow_indices, col_indices, values, size, dtype, layout, device, pin_memory);
  12624. }
  12625. auto [crow_indices_value, crow_indices_bdim] = unwrapTensorAtLevel(crow_indices, cur_level);
  12626. auto [col_indices_value, col_indices_bdim] = unwrapTensorAtLevel(col_indices, cur_level);
  12627. auto [values_value, values_bdim] = unwrapTensorAtLevel(values, cur_level);
  12628. auto results = batch_rule(crow_indices_value, crow_indices_bdim, col_indices_value, col_indices_bdim, values_value, values_bdim, size, dtype, layout, device, pin_memory);
  12629. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12630. }
  12631. template <typename batch_rule_t, batch_rule_t batch_rule>
  12632. at::Tensor sparse_csc_tensor_ccol_row_value_size_generated_plumbing(const at::Tensor & ccol_indices, const at::Tensor & row_indices, const at::Tensor & values, at::IntArrayRef size, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory) {
  12633. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12634. auto maybe_layer = maybeCurrentDynamicLayer();
  12635. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12636. int64_t cur_level = maybe_layer->layerId();
  12637. if (!isBatchedAtLevel(ccol_indices, cur_level) && !isBatchedAtLevel(row_indices, cur_level) && !isBatchedAtLevel(values, cur_level)) {
  12638. return at::_ops::sparse_csc_tensor_ccol_row_value_size::call(ccol_indices, row_indices, values, size, dtype, layout, device, pin_memory);
  12639. }
  12640. auto [ccol_indices_value, ccol_indices_bdim] = unwrapTensorAtLevel(ccol_indices, cur_level);
  12641. auto [row_indices_value, row_indices_bdim] = unwrapTensorAtLevel(row_indices, cur_level);
  12642. auto [values_value, values_bdim] = unwrapTensorAtLevel(values, cur_level);
  12643. auto results = batch_rule(ccol_indices_value, ccol_indices_bdim, row_indices_value, row_indices_bdim, values_value, values_bdim, size, dtype, layout, device, pin_memory);
  12644. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12645. }
  12646. template <typename batch_rule_t, batch_rule_t batch_rule>
  12647. at::Tensor sparse_bsr_tensor_crow_col_value_size_generated_plumbing(const at::Tensor & crow_indices, const at::Tensor & col_indices, const at::Tensor & values, at::IntArrayRef size, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory) {
  12648. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12649. auto maybe_layer = maybeCurrentDynamicLayer();
  12650. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12651. int64_t cur_level = maybe_layer->layerId();
  12652. if (!isBatchedAtLevel(crow_indices, cur_level) && !isBatchedAtLevel(col_indices, cur_level) && !isBatchedAtLevel(values, cur_level)) {
  12653. return at::_ops::sparse_bsr_tensor_crow_col_value_size::call(crow_indices, col_indices, values, size, dtype, layout, device, pin_memory);
  12654. }
  12655. auto [crow_indices_value, crow_indices_bdim] = unwrapTensorAtLevel(crow_indices, cur_level);
  12656. auto [col_indices_value, col_indices_bdim] = unwrapTensorAtLevel(col_indices, cur_level);
  12657. auto [values_value, values_bdim] = unwrapTensorAtLevel(values, cur_level);
  12658. auto results = batch_rule(crow_indices_value, crow_indices_bdim, col_indices_value, col_indices_bdim, values_value, values_bdim, size, dtype, layout, device, pin_memory);
  12659. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12660. }
  12661. template <typename batch_rule_t, batch_rule_t batch_rule>
  12662. at::Tensor sparse_bsc_tensor_ccol_row_value_size_generated_plumbing(const at::Tensor & ccol_indices, const at::Tensor & row_indices, const at::Tensor & values, at::IntArrayRef size, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory) {
  12663. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12664. auto maybe_layer = maybeCurrentDynamicLayer();
  12665. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12666. int64_t cur_level = maybe_layer->layerId();
  12667. if (!isBatchedAtLevel(ccol_indices, cur_level) && !isBatchedAtLevel(row_indices, cur_level) && !isBatchedAtLevel(values, cur_level)) {
  12668. return at::_ops::sparse_bsc_tensor_ccol_row_value_size::call(ccol_indices, row_indices, values, size, dtype, layout, device, pin_memory);
  12669. }
  12670. auto [ccol_indices_value, ccol_indices_bdim] = unwrapTensorAtLevel(ccol_indices, cur_level);
  12671. auto [row_indices_value, row_indices_bdim] = unwrapTensorAtLevel(row_indices, cur_level);
  12672. auto [values_value, values_bdim] = unwrapTensorAtLevel(values, cur_level);
  12673. auto results = batch_rule(ccol_indices_value, ccol_indices_bdim, row_indices_value, row_indices_bdim, values_value, values_bdim, size, dtype, layout, device, pin_memory);
  12674. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12675. }
  12676. template <typename batch_rule_t, batch_rule_t batch_rule>
  12677. at::Tensor sparse_compressed_tensor_comp_plain_value_generated_plumbing(const at::Tensor & compressed_indices, const at::Tensor & plain_indices, const at::Tensor & values, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory) {
  12678. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12679. auto maybe_layer = maybeCurrentDynamicLayer();
  12680. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12681. int64_t cur_level = maybe_layer->layerId();
  12682. if (!isBatchedAtLevel(compressed_indices, cur_level) && !isBatchedAtLevel(plain_indices, cur_level) && !isBatchedAtLevel(values, cur_level)) {
  12683. return at::_ops::sparse_compressed_tensor_comp_plain_value::call(compressed_indices, plain_indices, values, dtype, layout, device, pin_memory);
  12684. }
  12685. auto [compressed_indices_value, compressed_indices_bdim] = unwrapTensorAtLevel(compressed_indices, cur_level);
  12686. auto [plain_indices_value, plain_indices_bdim] = unwrapTensorAtLevel(plain_indices, cur_level);
  12687. auto [values_value, values_bdim] = unwrapTensorAtLevel(values, cur_level);
  12688. auto results = batch_rule(compressed_indices_value, compressed_indices_bdim, plain_indices_value, plain_indices_bdim, values_value, values_bdim, dtype, layout, device, pin_memory);
  12689. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12690. }
  12691. template <typename batch_rule_t, batch_rule_t batch_rule>
  12692. at::Tensor sparse_csr_tensor_crow_col_value_generated_plumbing(const at::Tensor & crow_indices, const at::Tensor & col_indices, const at::Tensor & values, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory) {
  12693. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12694. auto maybe_layer = maybeCurrentDynamicLayer();
  12695. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12696. int64_t cur_level = maybe_layer->layerId();
  12697. if (!isBatchedAtLevel(crow_indices, cur_level) && !isBatchedAtLevel(col_indices, cur_level) && !isBatchedAtLevel(values, cur_level)) {
  12698. return at::_ops::sparse_csr_tensor_crow_col_value::call(crow_indices, col_indices, values, dtype, layout, device, pin_memory);
  12699. }
  12700. auto [crow_indices_value, crow_indices_bdim] = unwrapTensorAtLevel(crow_indices, cur_level);
  12701. auto [col_indices_value, col_indices_bdim] = unwrapTensorAtLevel(col_indices, cur_level);
  12702. auto [values_value, values_bdim] = unwrapTensorAtLevel(values, cur_level);
  12703. auto results = batch_rule(crow_indices_value, crow_indices_bdim, col_indices_value, col_indices_bdim, values_value, values_bdim, dtype, layout, device, pin_memory);
  12704. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12705. }
  12706. template <typename batch_rule_t, batch_rule_t batch_rule>
  12707. at::Tensor sparse_csc_tensor_ccol_row_value_generated_plumbing(const at::Tensor & ccol_indices, const at::Tensor & row_indices, const at::Tensor & values, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory) {
  12708. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12709. auto maybe_layer = maybeCurrentDynamicLayer();
  12710. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12711. int64_t cur_level = maybe_layer->layerId();
  12712. if (!isBatchedAtLevel(ccol_indices, cur_level) && !isBatchedAtLevel(row_indices, cur_level) && !isBatchedAtLevel(values, cur_level)) {
  12713. return at::_ops::sparse_csc_tensor_ccol_row_value::call(ccol_indices, row_indices, values, dtype, layout, device, pin_memory);
  12714. }
  12715. auto [ccol_indices_value, ccol_indices_bdim] = unwrapTensorAtLevel(ccol_indices, cur_level);
  12716. auto [row_indices_value, row_indices_bdim] = unwrapTensorAtLevel(row_indices, cur_level);
  12717. auto [values_value, values_bdim] = unwrapTensorAtLevel(values, cur_level);
  12718. auto results = batch_rule(ccol_indices_value, ccol_indices_bdim, row_indices_value, row_indices_bdim, values_value, values_bdim, dtype, layout, device, pin_memory);
  12719. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12720. }
  12721. template <typename batch_rule_t, batch_rule_t batch_rule>
  12722. at::Tensor sparse_bsr_tensor_crow_col_value_generated_plumbing(const at::Tensor & crow_indices, const at::Tensor & col_indices, const at::Tensor & values, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory) {
  12723. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12724. auto maybe_layer = maybeCurrentDynamicLayer();
  12725. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12726. int64_t cur_level = maybe_layer->layerId();
  12727. if (!isBatchedAtLevel(crow_indices, cur_level) && !isBatchedAtLevel(col_indices, cur_level) && !isBatchedAtLevel(values, cur_level)) {
  12728. return at::_ops::sparse_bsr_tensor_crow_col_value::call(crow_indices, col_indices, values, dtype, layout, device, pin_memory);
  12729. }
  12730. auto [crow_indices_value, crow_indices_bdim] = unwrapTensorAtLevel(crow_indices, cur_level);
  12731. auto [col_indices_value, col_indices_bdim] = unwrapTensorAtLevel(col_indices, cur_level);
  12732. auto [values_value, values_bdim] = unwrapTensorAtLevel(values, cur_level);
  12733. auto results = batch_rule(crow_indices_value, crow_indices_bdim, col_indices_value, col_indices_bdim, values_value, values_bdim, dtype, layout, device, pin_memory);
  12734. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12735. }
  12736. template <typename batch_rule_t, batch_rule_t batch_rule>
  12737. at::Tensor sparse_bsc_tensor_ccol_row_value_generated_plumbing(const at::Tensor & ccol_indices, const at::Tensor & row_indices, const at::Tensor & values, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory) {
  12738. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12739. auto maybe_layer = maybeCurrentDynamicLayer();
  12740. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12741. int64_t cur_level = maybe_layer->layerId();
  12742. if (!isBatchedAtLevel(ccol_indices, cur_level) && !isBatchedAtLevel(row_indices, cur_level) && !isBatchedAtLevel(values, cur_level)) {
  12743. return at::_ops::sparse_bsc_tensor_ccol_row_value::call(ccol_indices, row_indices, values, dtype, layout, device, pin_memory);
  12744. }
  12745. auto [ccol_indices_value, ccol_indices_bdim] = unwrapTensorAtLevel(ccol_indices, cur_level);
  12746. auto [row_indices_value, row_indices_bdim] = unwrapTensorAtLevel(row_indices, cur_level);
  12747. auto [values_value, values_bdim] = unwrapTensorAtLevel(values, cur_level);
  12748. auto results = batch_rule(ccol_indices_value, ccol_indices_bdim, row_indices_value, row_indices_bdim, values_value, values_bdim, dtype, layout, device, pin_memory);
  12749. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12750. }
  12751. template <typename batch_rule_t, batch_rule_t batch_rule>
  12752. at::Tensor _sparse_compressed_tensor_unsafe_generated_plumbing(const at::Tensor & compressed_indices, const at::Tensor & plain_indices, const at::Tensor & values, c10::SymIntArrayRef size, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory) {
  12753. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12754. auto maybe_layer = maybeCurrentDynamicLayer();
  12755. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12756. int64_t cur_level = maybe_layer->layerId();
  12757. if (!isBatchedAtLevel(compressed_indices, cur_level) && !isBatchedAtLevel(plain_indices, cur_level) && !isBatchedAtLevel(values, cur_level)) {
  12758. return at::_ops::_sparse_compressed_tensor_unsafe::call(compressed_indices, plain_indices, values, size, dtype, layout, device, pin_memory);
  12759. }
  12760. auto [compressed_indices_value, compressed_indices_bdim] = unwrapTensorAtLevel(compressed_indices, cur_level);
  12761. auto [plain_indices_value, plain_indices_bdim] = unwrapTensorAtLevel(plain_indices, cur_level);
  12762. auto [values_value, values_bdim] = unwrapTensorAtLevel(values, cur_level);
  12763. auto results = batch_rule(compressed_indices_value, compressed_indices_bdim, plain_indices_value, plain_indices_bdim, values_value, values_bdim, size, dtype, layout, device, pin_memory);
  12764. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12765. }
  12766. template <typename batch_rule_t, batch_rule_t batch_rule>
  12767. at::Tensor _sparse_csr_tensor_unsafe_generated_plumbing(const at::Tensor & crow_indices, const at::Tensor & col_indices, const at::Tensor & values, at::IntArrayRef size, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory) {
  12768. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12769. auto maybe_layer = maybeCurrentDynamicLayer();
  12770. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12771. int64_t cur_level = maybe_layer->layerId();
  12772. if (!isBatchedAtLevel(crow_indices, cur_level) && !isBatchedAtLevel(col_indices, cur_level) && !isBatchedAtLevel(values, cur_level)) {
  12773. return at::_ops::_sparse_csr_tensor_unsafe::call(crow_indices, col_indices, values, size, dtype, layout, device, pin_memory);
  12774. }
  12775. auto [crow_indices_value, crow_indices_bdim] = unwrapTensorAtLevel(crow_indices, cur_level);
  12776. auto [col_indices_value, col_indices_bdim] = unwrapTensorAtLevel(col_indices, cur_level);
  12777. auto [values_value, values_bdim] = unwrapTensorAtLevel(values, cur_level);
  12778. auto results = batch_rule(crow_indices_value, crow_indices_bdim, col_indices_value, col_indices_bdim, values_value, values_bdim, size, dtype, layout, device, pin_memory);
  12779. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12780. }
  12781. template <typename batch_rule_t, batch_rule_t batch_rule>
  12782. at::Tensor _sparse_csc_tensor_unsafe_generated_plumbing(const at::Tensor & ccol_indices, const at::Tensor & row_indices, const at::Tensor & values, at::IntArrayRef size, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory) {
  12783. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12784. auto maybe_layer = maybeCurrentDynamicLayer();
  12785. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12786. int64_t cur_level = maybe_layer->layerId();
  12787. if (!isBatchedAtLevel(ccol_indices, cur_level) && !isBatchedAtLevel(row_indices, cur_level) && !isBatchedAtLevel(values, cur_level)) {
  12788. return at::_ops::_sparse_csc_tensor_unsafe::call(ccol_indices, row_indices, values, size, dtype, layout, device, pin_memory);
  12789. }
  12790. auto [ccol_indices_value, ccol_indices_bdim] = unwrapTensorAtLevel(ccol_indices, cur_level);
  12791. auto [row_indices_value, row_indices_bdim] = unwrapTensorAtLevel(row_indices, cur_level);
  12792. auto [values_value, values_bdim] = unwrapTensorAtLevel(values, cur_level);
  12793. auto results = batch_rule(ccol_indices_value, ccol_indices_bdim, row_indices_value, row_indices_bdim, values_value, values_bdim, size, dtype, layout, device, pin_memory);
  12794. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12795. }
  12796. template <typename batch_rule_t, batch_rule_t batch_rule>
  12797. at::Tensor _sparse_bsr_tensor_unsafe_generated_plumbing(const at::Tensor & crow_indices, const at::Tensor & col_indices, const at::Tensor & values, at::IntArrayRef size, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory) {
  12798. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12799. auto maybe_layer = maybeCurrentDynamicLayer();
  12800. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12801. int64_t cur_level = maybe_layer->layerId();
  12802. if (!isBatchedAtLevel(crow_indices, cur_level) && !isBatchedAtLevel(col_indices, cur_level) && !isBatchedAtLevel(values, cur_level)) {
  12803. return at::_ops::_sparse_bsr_tensor_unsafe::call(crow_indices, col_indices, values, size, dtype, layout, device, pin_memory);
  12804. }
  12805. auto [crow_indices_value, crow_indices_bdim] = unwrapTensorAtLevel(crow_indices, cur_level);
  12806. auto [col_indices_value, col_indices_bdim] = unwrapTensorAtLevel(col_indices, cur_level);
  12807. auto [values_value, values_bdim] = unwrapTensorAtLevel(values, cur_level);
  12808. auto results = batch_rule(crow_indices_value, crow_indices_bdim, col_indices_value, col_indices_bdim, values_value, values_bdim, size, dtype, layout, device, pin_memory);
  12809. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12810. }
  12811. template <typename batch_rule_t, batch_rule_t batch_rule>
  12812. at::Tensor _sparse_bsc_tensor_unsafe_generated_plumbing(const at::Tensor & ccol_indices, const at::Tensor & row_indices, const at::Tensor & values, at::IntArrayRef size, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory) {
  12813. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12814. auto maybe_layer = maybeCurrentDynamicLayer();
  12815. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12816. int64_t cur_level = maybe_layer->layerId();
  12817. if (!isBatchedAtLevel(ccol_indices, cur_level) && !isBatchedAtLevel(row_indices, cur_level) && !isBatchedAtLevel(values, cur_level)) {
  12818. return at::_ops::_sparse_bsc_tensor_unsafe::call(ccol_indices, row_indices, values, size, dtype, layout, device, pin_memory);
  12819. }
  12820. auto [ccol_indices_value, ccol_indices_bdim] = unwrapTensorAtLevel(ccol_indices, cur_level);
  12821. auto [row_indices_value, row_indices_bdim] = unwrapTensorAtLevel(row_indices, cur_level);
  12822. auto [values_value, values_bdim] = unwrapTensorAtLevel(values, cur_level);
  12823. auto results = batch_rule(ccol_indices_value, ccol_indices_bdim, row_indices_value, row_indices_bdim, values_value, values_bdim, size, dtype, layout, device, pin_memory);
  12824. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12825. }
  12826. template <typename batch_rule_t, batch_rule_t batch_rule>
  12827. at::Tensor sparse_coo_tensor_indices_generated_plumbing(const at::Tensor & indices, const at::Tensor & values, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory, ::std::optional<bool> is_coalesced) {
  12828. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12829. auto maybe_layer = maybeCurrentDynamicLayer();
  12830. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12831. int64_t cur_level = maybe_layer->layerId();
  12832. if (!isBatchedAtLevel(indices, cur_level) && !isBatchedAtLevel(values, cur_level)) {
  12833. return at::_ops::sparse_coo_tensor_indices::call(indices, values, dtype, layout, device, pin_memory, is_coalesced);
  12834. }
  12835. auto [indices_value, indices_bdim] = unwrapTensorAtLevel(indices, cur_level);
  12836. auto [values_value, values_bdim] = unwrapTensorAtLevel(values, cur_level);
  12837. auto results = batch_rule(indices_value, indices_bdim, values_value, values_bdim, dtype, layout, device, pin_memory, is_coalesced);
  12838. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12839. }
  12840. template <typename batch_rule_t, batch_rule_t batch_rule>
  12841. at::Tensor sparse_coo_tensor_indices_size_generated_plumbing(const at::Tensor & indices, const at::Tensor & values, at::IntArrayRef size, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory, ::std::optional<bool> is_coalesced) {
  12842. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12843. auto maybe_layer = maybeCurrentDynamicLayer();
  12844. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12845. int64_t cur_level = maybe_layer->layerId();
  12846. if (!isBatchedAtLevel(indices, cur_level) && !isBatchedAtLevel(values, cur_level)) {
  12847. return at::_ops::sparse_coo_tensor_indices_size::call(indices, values, size, dtype, layout, device, pin_memory, is_coalesced);
  12848. }
  12849. auto [indices_value, indices_bdim] = unwrapTensorAtLevel(indices, cur_level);
  12850. auto [values_value, values_bdim] = unwrapTensorAtLevel(values, cur_level);
  12851. auto results = batch_rule(indices_value, indices_bdim, values_value, values_bdim, size, dtype, layout, device, pin_memory, is_coalesced);
  12852. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12853. }
  12854. template <typename batch_rule_t, batch_rule_t batch_rule>
  12855. at::Tensor _sparse_coo_tensor_unsafe_generated_plumbing(const at::Tensor & indices, const at::Tensor & values, c10::SymIntArrayRef size, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory, ::std::optional<bool> is_coalesced) {
  12856. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12857. auto maybe_layer = maybeCurrentDynamicLayer();
  12858. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12859. int64_t cur_level = maybe_layer->layerId();
  12860. if (!isBatchedAtLevel(indices, cur_level) && !isBatchedAtLevel(values, cur_level)) {
  12861. return at::_ops::_sparse_coo_tensor_unsafe::call(indices, values, size, dtype, layout, device, pin_memory, is_coalesced);
  12862. }
  12863. auto [indices_value, indices_bdim] = unwrapTensorAtLevel(indices, cur_level);
  12864. auto [values_value, values_bdim] = unwrapTensorAtLevel(values, cur_level);
  12865. auto results = batch_rule(indices_value, indices_bdim, values_value, values_bdim, size, dtype, layout, device, pin_memory, is_coalesced);
  12866. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12867. }
  12868. template <typename batch_rule_t, batch_rule_t batch_rule>
  12869. void _validate_sparse_coo_tensor_args_generated_plumbing(const at::Tensor & indices, const at::Tensor & values, at::IntArrayRef size, ::std::optional<bool> is_coalesced, ::std::optional<bool> check_pinning) {
  12870. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12871. auto maybe_layer = maybeCurrentDynamicLayer();
  12872. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  12873. int64_t cur_level = maybe_layer->layerId();
  12874. if (!isBatchedAtLevel(indices, cur_level) && !isBatchedAtLevel(values, cur_level)) {
  12875. return at::_ops::_validate_sparse_coo_tensor_args::call(indices, values, size, is_coalesced, check_pinning);
  12876. }
  12877. auto [indices_value, indices_bdim] = unwrapTensorAtLevel(indices, cur_level);
  12878. auto [values_value, values_bdim] = unwrapTensorAtLevel(values, cur_level);
  12879. batch_rule(indices_value, indices_bdim, values_value, values_bdim, size, is_coalesced, check_pinning);
  12880. }
  12881. template <typename batch_rule_t, batch_rule_t batch_rule>
  12882. void _validate_sparse_compressed_tensor_args_generated_plumbing(const at::Tensor & compressed_indices, const at::Tensor & plain_indices, const at::Tensor & values, at::IntArrayRef size, at::Layout layout, ::std::optional<bool> check_pinning) {
  12883. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12884. auto maybe_layer = maybeCurrentDynamicLayer();
  12885. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  12886. int64_t cur_level = maybe_layer->layerId();
  12887. if (!isBatchedAtLevel(compressed_indices, cur_level) && !isBatchedAtLevel(plain_indices, cur_level) && !isBatchedAtLevel(values, cur_level)) {
  12888. return at::_ops::_validate_sparse_compressed_tensor_args::call(compressed_indices, plain_indices, values, size, layout, check_pinning);
  12889. }
  12890. auto [compressed_indices_value, compressed_indices_bdim] = unwrapTensorAtLevel(compressed_indices, cur_level);
  12891. auto [plain_indices_value, plain_indices_bdim] = unwrapTensorAtLevel(plain_indices, cur_level);
  12892. auto [values_value, values_bdim] = unwrapTensorAtLevel(values, cur_level);
  12893. batch_rule(compressed_indices_value, compressed_indices_bdim, plain_indices_value, plain_indices_bdim, values_value, values_bdim, size, layout, check_pinning);
  12894. }
  12895. template <typename batch_rule_t, batch_rule_t batch_rule>
  12896. void _validate_sparse_csr_tensor_args_generated_plumbing(const at::Tensor & crow_indices, const at::Tensor & col_indices, const at::Tensor & values, at::IntArrayRef size, ::std::optional<bool> check_pinning) {
  12897. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12898. auto maybe_layer = maybeCurrentDynamicLayer();
  12899. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  12900. int64_t cur_level = maybe_layer->layerId();
  12901. if (!isBatchedAtLevel(crow_indices, cur_level) && !isBatchedAtLevel(col_indices, cur_level) && !isBatchedAtLevel(values, cur_level)) {
  12902. return at::_ops::_validate_sparse_csr_tensor_args::call(crow_indices, col_indices, values, size, check_pinning);
  12903. }
  12904. auto [crow_indices_value, crow_indices_bdim] = unwrapTensorAtLevel(crow_indices, cur_level);
  12905. auto [col_indices_value, col_indices_bdim] = unwrapTensorAtLevel(col_indices, cur_level);
  12906. auto [values_value, values_bdim] = unwrapTensorAtLevel(values, cur_level);
  12907. batch_rule(crow_indices_value, crow_indices_bdim, col_indices_value, col_indices_bdim, values_value, values_bdim, size, check_pinning);
  12908. }
  12909. template <typename batch_rule_t, batch_rule_t batch_rule>
  12910. void _validate_sparse_csc_tensor_args_generated_plumbing(const at::Tensor & ccol_indices, const at::Tensor & row_indices, const at::Tensor & values, at::IntArrayRef size, ::std::optional<bool> check_pinning) {
  12911. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12912. auto maybe_layer = maybeCurrentDynamicLayer();
  12913. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  12914. int64_t cur_level = maybe_layer->layerId();
  12915. if (!isBatchedAtLevel(ccol_indices, cur_level) && !isBatchedAtLevel(row_indices, cur_level) && !isBatchedAtLevel(values, cur_level)) {
  12916. return at::_ops::_validate_sparse_csc_tensor_args::call(ccol_indices, row_indices, values, size, check_pinning);
  12917. }
  12918. auto [ccol_indices_value, ccol_indices_bdim] = unwrapTensorAtLevel(ccol_indices, cur_level);
  12919. auto [row_indices_value, row_indices_bdim] = unwrapTensorAtLevel(row_indices, cur_level);
  12920. auto [values_value, values_bdim] = unwrapTensorAtLevel(values, cur_level);
  12921. batch_rule(ccol_indices_value, ccol_indices_bdim, row_indices_value, row_indices_bdim, values_value, values_bdim, size, check_pinning);
  12922. }
  12923. template <typename batch_rule_t, batch_rule_t batch_rule>
  12924. void _validate_sparse_bsr_tensor_args_generated_plumbing(const at::Tensor & crow_indices, const at::Tensor & col_indices, const at::Tensor & values, at::IntArrayRef size, ::std::optional<bool> check_pinning) {
  12925. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12926. auto maybe_layer = maybeCurrentDynamicLayer();
  12927. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  12928. int64_t cur_level = maybe_layer->layerId();
  12929. if (!isBatchedAtLevel(crow_indices, cur_level) && !isBatchedAtLevel(col_indices, cur_level) && !isBatchedAtLevel(values, cur_level)) {
  12930. return at::_ops::_validate_sparse_bsr_tensor_args::call(crow_indices, col_indices, values, size, check_pinning);
  12931. }
  12932. auto [crow_indices_value, crow_indices_bdim] = unwrapTensorAtLevel(crow_indices, cur_level);
  12933. auto [col_indices_value, col_indices_bdim] = unwrapTensorAtLevel(col_indices, cur_level);
  12934. auto [values_value, values_bdim] = unwrapTensorAtLevel(values, cur_level);
  12935. batch_rule(crow_indices_value, crow_indices_bdim, col_indices_value, col_indices_bdim, values_value, values_bdim, size, check_pinning);
  12936. }
  12937. template <typename batch_rule_t, batch_rule_t batch_rule>
  12938. void _validate_sparse_bsc_tensor_args_generated_plumbing(const at::Tensor & ccol_indices, const at::Tensor & row_indices, const at::Tensor & values, at::IntArrayRef size, ::std::optional<bool> check_pinning) {
  12939. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12940. auto maybe_layer = maybeCurrentDynamicLayer();
  12941. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  12942. int64_t cur_level = maybe_layer->layerId();
  12943. if (!isBatchedAtLevel(ccol_indices, cur_level) && !isBatchedAtLevel(row_indices, cur_level) && !isBatchedAtLevel(values, cur_level)) {
  12944. return at::_ops::_validate_sparse_bsc_tensor_args::call(ccol_indices, row_indices, values, size, check_pinning);
  12945. }
  12946. auto [ccol_indices_value, ccol_indices_bdim] = unwrapTensorAtLevel(ccol_indices, cur_level);
  12947. auto [row_indices_value, row_indices_bdim] = unwrapTensorAtLevel(row_indices, cur_level);
  12948. auto [values_value, values_bdim] = unwrapTensorAtLevel(values, cur_level);
  12949. batch_rule(ccol_indices_value, ccol_indices_bdim, row_indices_value, row_indices_bdim, values_value, values_bdim, size, check_pinning);
  12950. }
  12951. template <typename batch_rule_t, batch_rule_t batch_rule>
  12952. at::Tensor _sparse_coo_tensor_with_dims_and_tensors_generated_plumbing(int64_t sparse_dim, int64_t dense_dim, c10::SymIntArrayRef size, const at::Tensor & indices, const at::Tensor & values, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory, ::std::optional<bool> is_coalesced) {
  12953. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12954. auto maybe_layer = maybeCurrentDynamicLayer();
  12955. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12956. int64_t cur_level = maybe_layer->layerId();
  12957. if (!isBatchedAtLevel(indices, cur_level) && !isBatchedAtLevel(values, cur_level)) {
  12958. return at::_ops::_sparse_coo_tensor_with_dims_and_tensors::call(sparse_dim, dense_dim, size, indices, values, dtype, layout, device, pin_memory, is_coalesced);
  12959. }
  12960. auto [indices_value, indices_bdim] = unwrapTensorAtLevel(indices, cur_level);
  12961. auto [values_value, values_bdim] = unwrapTensorAtLevel(values, cur_level);
  12962. auto results = batch_rule(sparse_dim, dense_dim, size, indices_value, indices_bdim, values_value, values_bdim, dtype, layout, device, pin_memory, is_coalesced);
  12963. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  12964. }
  12965. template <typename batch_rule_t, batch_rule_t batch_rule>
  12966. const at::Tensor & sparse_resize__generated_plumbing(const at::Tensor & self, at::IntArrayRef size, int64_t sparse_dim, int64_t dense_dim) {
  12967. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12968. auto maybe_layer = maybeCurrentDynamicLayer();
  12969. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  12970. int64_t cur_level = maybe_layer->layerId();
  12971. if (!isBatchedAtLevel(self, cur_level)) {
  12972. return at::_ops::sparse_resize_::call(self, size, sparse_dim, dense_dim);
  12973. }
  12974. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  12975. batch_rule(self_value, self_bdim, size, sparse_dim, dense_dim);
  12976. return self;
  12977. }
  12978. template <typename batch_rule_t, batch_rule_t batch_rule>
  12979. const at::Tensor & sparse_resize_and_clear__generated_plumbing(const at::Tensor & self, at::IntArrayRef size, int64_t sparse_dim, int64_t dense_dim) {
  12980. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12981. auto maybe_layer = maybeCurrentDynamicLayer();
  12982. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  12983. int64_t cur_level = maybe_layer->layerId();
  12984. if (!isBatchedAtLevel(self, cur_level)) {
  12985. return at::_ops::sparse_resize_and_clear_::call(self, size, sparse_dim, dense_dim);
  12986. }
  12987. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  12988. batch_rule(self_value, self_bdim, size, sparse_dim, dense_dim);
  12989. return self;
  12990. }
  12991. template <typename batch_rule_t, batch_rule_t batch_rule>
  12992. at::Tensor sparse_mask_generated_plumbing(const at::Tensor & self, const at::Tensor & mask) {
  12993. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  12994. auto maybe_layer = maybeCurrentDynamicLayer();
  12995. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  12996. int64_t cur_level = maybe_layer->layerId();
  12997. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(mask, cur_level)) {
  12998. return at::_ops::sparse_mask::call(self, mask);
  12999. }
  13000. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13001. auto [mask_value, mask_bdim] = unwrapTensorAtLevel(mask, cur_level);
  13002. auto results = batch_rule(self_value, self_bdim, mask_value, mask_bdim);
  13003. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13004. }
  13005. template <typename batch_rule_t, batch_rule_t batch_rule>
  13006. at::Tensor _sparse_mask_projection_generated_plumbing(const at::Tensor & self, const at::Tensor & mask, bool accumulate_matches) {
  13007. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13008. auto maybe_layer = maybeCurrentDynamicLayer();
  13009. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13010. int64_t cur_level = maybe_layer->layerId();
  13011. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(mask, cur_level)) {
  13012. return at::_ops::_sparse_mask_projection::call(self, mask, accumulate_matches);
  13013. }
  13014. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13015. auto [mask_value, mask_bdim] = unwrapTensorAtLevel(mask, cur_level);
  13016. auto results = batch_rule(self_value, self_bdim, mask_value, mask_bdim, accumulate_matches);
  13017. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13018. }
  13019. template <typename batch_rule_t, batch_rule_t batch_rule>
  13020. ::std::vector<at::Tensor> _to_cpu_generated_plumbing(at::TensorList tensors) {
  13021. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13022. auto maybe_layer = maybeCurrentDynamicLayer();
  13023. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13024. int64_t cur_level = maybe_layer->layerId();
  13025. if (!isBatchedAtLevel(tensors, cur_level)) {
  13026. return at::_ops::_to_cpu::call(tensors);
  13027. }
  13028. auto results = batch_rule(tensors);
  13029. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  13030. }
  13031. template <typename batch_rule_t, batch_rule_t batch_rule>
  13032. at::Tensor to_dense_generated_plumbing(const at::Tensor & self, ::std::optional<at::ScalarType> dtype, ::std::optional<bool> masked_grad) {
  13033. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13034. auto maybe_layer = maybeCurrentDynamicLayer();
  13035. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13036. int64_t cur_level = maybe_layer->layerId();
  13037. if (!isBatchedAtLevel(self, cur_level)) {
  13038. return at::_ops::to_dense::call(self, dtype, masked_grad);
  13039. }
  13040. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13041. auto results = batch_rule(self_value, self_bdim, dtype, masked_grad);
  13042. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13043. }
  13044. template <typename batch_rule_t, batch_rule_t batch_rule>
  13045. at::Tensor _to_dense_generated_plumbing(const at::Tensor & self, ::std::optional<at::ScalarType> dtype, ::std::optional<bool> masked_grad) {
  13046. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13047. auto maybe_layer = maybeCurrentDynamicLayer();
  13048. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13049. int64_t cur_level = maybe_layer->layerId();
  13050. if (!isBatchedAtLevel(self, cur_level)) {
  13051. return at::_ops::_to_dense::call(self, dtype, masked_grad);
  13052. }
  13053. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13054. auto results = batch_rule(self_value, self_bdim, dtype, masked_grad);
  13055. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13056. }
  13057. template <typename batch_rule_t, batch_rule_t batch_rule>
  13058. at::Tensor to_dense_backward_generated_plumbing(const at::Tensor & grad, const at::Tensor & input, ::std::optional<bool> masked_grad) {
  13059. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13060. auto maybe_layer = maybeCurrentDynamicLayer();
  13061. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13062. int64_t cur_level = maybe_layer->layerId();
  13063. if (!isBatchedAtLevel(grad, cur_level) && !isBatchedAtLevel(input, cur_level)) {
  13064. return at::_ops::to_dense_backward::call(grad, input, masked_grad);
  13065. }
  13066. auto [grad_value, grad_bdim] = unwrapTensorAtLevel(grad, cur_level);
  13067. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  13068. auto results = batch_rule(grad_value, grad_bdim, input_value, input_bdim, masked_grad);
  13069. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13070. }
  13071. template <typename batch_rule_t, batch_rule_t batch_rule>
  13072. at::Tensor coalesce_generated_plumbing(const at::Tensor & self) {
  13073. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13074. auto maybe_layer = maybeCurrentDynamicLayer();
  13075. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13076. int64_t cur_level = maybe_layer->layerId();
  13077. if (!isBatchedAtLevel(self, cur_level)) {
  13078. return at::_ops::coalesce::call(self);
  13079. }
  13080. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13081. auto results = batch_rule(self_value, self_bdim);
  13082. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13083. }
  13084. template <typename batch_rule_t, batch_rule_t batch_rule>
  13085. at::Tensor _coalesce_generated_plumbing(const at::Tensor & self) {
  13086. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13087. auto maybe_layer = maybeCurrentDynamicLayer();
  13088. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13089. int64_t cur_level = maybe_layer->layerId();
  13090. if (!isBatchedAtLevel(self, cur_level)) {
  13091. return at::_ops::_coalesce::call(self);
  13092. }
  13093. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13094. auto results = batch_rule(self_value, self_bdim);
  13095. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13096. }
  13097. template <typename batch_rule_t, batch_rule_t batch_rule>
  13098. at::Tensor _indices_generated_plumbing(const at::Tensor & self) {
  13099. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13100. auto maybe_layer = maybeCurrentDynamicLayer();
  13101. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13102. int64_t cur_level = maybe_layer->layerId();
  13103. if (!isBatchedAtLevel(self, cur_level)) {
  13104. return at::_ops::_indices::call(self);
  13105. }
  13106. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13107. auto results = batch_rule(self_value, self_bdim);
  13108. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13109. }
  13110. template <typename batch_rule_t, batch_rule_t batch_rule>
  13111. at::Tensor _values_generated_plumbing(const at::Tensor & self) {
  13112. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13113. auto maybe_layer = maybeCurrentDynamicLayer();
  13114. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13115. int64_t cur_level = maybe_layer->layerId();
  13116. if (!isBatchedAtLevel(self, cur_level)) {
  13117. return at::_ops::_values::call(self);
  13118. }
  13119. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13120. auto results = batch_rule(self_value, self_bdim);
  13121. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13122. }
  13123. template <typename batch_rule_t, batch_rule_t batch_rule>
  13124. at::Tensor & _coalesced__generated_plumbing(at::Tensor & self, bool coalesced) {
  13125. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13126. auto maybe_layer = maybeCurrentDynamicLayer();
  13127. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  13128. int64_t cur_level = maybe_layer->layerId();
  13129. if (!isBatchedAtLevel(self, cur_level)) {
  13130. return at::_ops::_coalesced_::call(self, coalesced);
  13131. }
  13132. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13133. batch_rule(self_value, self_bdim, coalesced);
  13134. return self;
  13135. }
  13136. template <typename batch_rule_t, batch_rule_t batch_rule>
  13137. at::Tensor indices_generated_plumbing(const at::Tensor & self) {
  13138. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13139. auto maybe_layer = maybeCurrentDynamicLayer();
  13140. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13141. int64_t cur_level = maybe_layer->layerId();
  13142. if (!isBatchedAtLevel(self, cur_level)) {
  13143. return at::_ops::indices::call(self);
  13144. }
  13145. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13146. auto results = batch_rule(self_value, self_bdim);
  13147. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13148. }
  13149. template <typename batch_rule_t, batch_rule_t batch_rule>
  13150. at::Tensor values_generated_plumbing(const at::Tensor & self) {
  13151. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13152. auto maybe_layer = maybeCurrentDynamicLayer();
  13153. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13154. int64_t cur_level = maybe_layer->layerId();
  13155. if (!isBatchedAtLevel(self, cur_level)) {
  13156. return at::_ops::values::call(self);
  13157. }
  13158. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13159. auto results = batch_rule(self_value, self_bdim);
  13160. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13161. }
  13162. template <typename batch_rule_t, batch_rule_t batch_rule>
  13163. at::Tensor crow_indices_generated_plumbing(const at::Tensor & self) {
  13164. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13165. auto maybe_layer = maybeCurrentDynamicLayer();
  13166. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13167. int64_t cur_level = maybe_layer->layerId();
  13168. if (!isBatchedAtLevel(self, cur_level)) {
  13169. return at::_ops::crow_indices::call(self);
  13170. }
  13171. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13172. auto results = batch_rule(self_value, self_bdim);
  13173. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13174. }
  13175. template <typename batch_rule_t, batch_rule_t batch_rule>
  13176. at::Tensor col_indices_generated_plumbing(const at::Tensor & self) {
  13177. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13178. auto maybe_layer = maybeCurrentDynamicLayer();
  13179. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13180. int64_t cur_level = maybe_layer->layerId();
  13181. if (!isBatchedAtLevel(self, cur_level)) {
  13182. return at::_ops::col_indices::call(self);
  13183. }
  13184. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13185. auto results = batch_rule(self_value, self_bdim);
  13186. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13187. }
  13188. template <typename batch_rule_t, batch_rule_t batch_rule>
  13189. at::Tensor ccol_indices_generated_plumbing(const at::Tensor & self) {
  13190. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13191. auto maybe_layer = maybeCurrentDynamicLayer();
  13192. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13193. int64_t cur_level = maybe_layer->layerId();
  13194. if (!isBatchedAtLevel(self, cur_level)) {
  13195. return at::_ops::ccol_indices::call(self);
  13196. }
  13197. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13198. auto results = batch_rule(self_value, self_bdim);
  13199. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13200. }
  13201. template <typename batch_rule_t, batch_rule_t batch_rule>
  13202. at::Tensor row_indices_generated_plumbing(const at::Tensor & self) {
  13203. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13204. auto maybe_layer = maybeCurrentDynamicLayer();
  13205. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13206. int64_t cur_level = maybe_layer->layerId();
  13207. if (!isBatchedAtLevel(self, cur_level)) {
  13208. return at::_ops::row_indices::call(self);
  13209. }
  13210. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13211. auto results = batch_rule(self_value, self_bdim);
  13212. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13213. }
  13214. template <typename batch_rule_t, batch_rule_t batch_rule>
  13215. at::Tensor hspmm_generated_plumbing(const at::Tensor & mat1, const at::Tensor & mat2) {
  13216. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13217. auto maybe_layer = maybeCurrentDynamicLayer();
  13218. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13219. int64_t cur_level = maybe_layer->layerId();
  13220. if (!isBatchedAtLevel(mat1, cur_level) && !isBatchedAtLevel(mat2, cur_level)) {
  13221. return at::_ops::hspmm::call(mat1, mat2);
  13222. }
  13223. auto [mat1_value, mat1_bdim] = unwrapTensorAtLevel(mat1, cur_level);
  13224. auto [mat2_value, mat2_bdim] = unwrapTensorAtLevel(mat2, cur_level);
  13225. auto results = batch_rule(mat1_value, mat1_bdim, mat2_value, mat2_bdim);
  13226. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13227. }
  13228. template <typename batch_rule_t, batch_rule_t batch_rule>
  13229. at::Tensor & copy_sparse_to_sparse__generated_plumbing(at::Tensor & self, const at::Tensor & src, bool non_blocking) {
  13230. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13231. auto maybe_layer = maybeCurrentDynamicLayer();
  13232. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  13233. int64_t cur_level = maybe_layer->layerId();
  13234. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(src, cur_level)) {
  13235. return at::_ops::copy_sparse_to_sparse_::call(self, src, non_blocking);
  13236. }
  13237. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13238. auto [src_value, src_bdim] = unwrapTensorAtLevel(src, cur_level);
  13239. batch_rule(self_value, self_bdim, src_value, src_bdim, non_blocking);
  13240. return self;
  13241. }
  13242. template <typename batch_rule_t, batch_rule_t batch_rule>
  13243. ::std::vector<at::Tensor> unbind_int_generated_plumbing(const at::Tensor & self, int64_t dim) {
  13244. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13245. auto maybe_layer = maybeCurrentDynamicLayer();
  13246. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13247. int64_t cur_level = maybe_layer->layerId();
  13248. if (!isBatchedAtLevel(self, cur_level)) {
  13249. return at::_ops::unbind_int::call(self, dim);
  13250. }
  13251. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13252. auto results = batch_rule(self_value, self_bdim, dim);
  13253. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  13254. }
  13255. template <typename batch_rule_t, batch_rule_t batch_rule>
  13256. ::std::vector<at::Tensor> unbind_Dimname_generated_plumbing(const at::Tensor & self, at::Dimname dim) {
  13257. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13258. auto maybe_layer = maybeCurrentDynamicLayer();
  13259. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13260. int64_t cur_level = maybe_layer->layerId();
  13261. if (!isBatchedAtLevel(self, cur_level)) {
  13262. return at::_ops::unbind_Dimname::call(self, dim);
  13263. }
  13264. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13265. auto results = batch_rule(self_value, self_bdim, dim);
  13266. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  13267. }
  13268. template <typename batch_rule_t, batch_rule_t batch_rule>
  13269. at::Tensor to_sparse_sparse_dim_generated_plumbing(const at::Tensor & self, int64_t sparse_dim) {
  13270. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13271. auto maybe_layer = maybeCurrentDynamicLayer();
  13272. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13273. int64_t cur_level = maybe_layer->layerId();
  13274. if (!isBatchedAtLevel(self, cur_level)) {
  13275. return at::_ops::to_sparse_sparse_dim::call(self, sparse_dim);
  13276. }
  13277. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13278. auto results = batch_rule(self_value, self_bdim, sparse_dim);
  13279. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13280. }
  13281. template <typename batch_rule_t, batch_rule_t batch_rule>
  13282. at::Tensor _to_sparse_sparse_dim_generated_plumbing(const at::Tensor & self, int64_t sparse_dim) {
  13283. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13284. auto maybe_layer = maybeCurrentDynamicLayer();
  13285. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13286. int64_t cur_level = maybe_layer->layerId();
  13287. if (!isBatchedAtLevel(self, cur_level)) {
  13288. return at::_ops::_to_sparse_sparse_dim::call(self, sparse_dim);
  13289. }
  13290. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13291. auto results = batch_rule(self_value, self_bdim, sparse_dim);
  13292. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13293. }
  13294. template <typename batch_rule_t, batch_rule_t batch_rule>
  13295. at::Tensor to_sparse_generated_plumbing(const at::Tensor & self, ::std::optional<at::Layout> layout, at::OptionalIntArrayRef blocksize, ::std::optional<int64_t> dense_dim) {
  13296. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13297. auto maybe_layer = maybeCurrentDynamicLayer();
  13298. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13299. int64_t cur_level = maybe_layer->layerId();
  13300. if (!isBatchedAtLevel(self, cur_level)) {
  13301. return at::_ops::to_sparse::call(self, layout, blocksize, dense_dim);
  13302. }
  13303. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13304. auto results = batch_rule(self_value, self_bdim, layout, blocksize, dense_dim);
  13305. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13306. }
  13307. template <typename batch_rule_t, batch_rule_t batch_rule>
  13308. at::Tensor _to_sparse_generated_plumbing(const at::Tensor & self, ::std::optional<at::Layout> layout, at::OptionalIntArrayRef blocksize, ::std::optional<int64_t> dense_dim) {
  13309. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13310. auto maybe_layer = maybeCurrentDynamicLayer();
  13311. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13312. int64_t cur_level = maybe_layer->layerId();
  13313. if (!isBatchedAtLevel(self, cur_level)) {
  13314. return at::_ops::_to_sparse::call(self, layout, blocksize, dense_dim);
  13315. }
  13316. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13317. auto results = batch_rule(self_value, self_bdim, layout, blocksize, dense_dim);
  13318. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13319. }
  13320. template <typename batch_rule_t, batch_rule_t batch_rule>
  13321. at::Tensor to_sparse_csr_generated_plumbing(const at::Tensor & self, ::std::optional<int64_t> dense_dim) {
  13322. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13323. auto maybe_layer = maybeCurrentDynamicLayer();
  13324. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13325. int64_t cur_level = maybe_layer->layerId();
  13326. if (!isBatchedAtLevel(self, cur_level)) {
  13327. return at::_ops::to_sparse_csr::call(self, dense_dim);
  13328. }
  13329. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13330. auto results = batch_rule(self_value, self_bdim, dense_dim);
  13331. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13332. }
  13333. template <typename batch_rule_t, batch_rule_t batch_rule>
  13334. at::Tensor _to_sparse_csr_generated_plumbing(const at::Tensor & self, ::std::optional<int64_t> dense_dim) {
  13335. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13336. auto maybe_layer = maybeCurrentDynamicLayer();
  13337. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13338. int64_t cur_level = maybe_layer->layerId();
  13339. if (!isBatchedAtLevel(self, cur_level)) {
  13340. return at::_ops::_to_sparse_csr::call(self, dense_dim);
  13341. }
  13342. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13343. auto results = batch_rule(self_value, self_bdim, dense_dim);
  13344. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13345. }
  13346. template <typename batch_rule_t, batch_rule_t batch_rule>
  13347. at::Tensor to_sparse_csc_generated_plumbing(const at::Tensor & self, ::std::optional<int64_t> dense_dim) {
  13348. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13349. auto maybe_layer = maybeCurrentDynamicLayer();
  13350. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13351. int64_t cur_level = maybe_layer->layerId();
  13352. if (!isBatchedAtLevel(self, cur_level)) {
  13353. return at::_ops::to_sparse_csc::call(self, dense_dim);
  13354. }
  13355. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13356. auto results = batch_rule(self_value, self_bdim, dense_dim);
  13357. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13358. }
  13359. template <typename batch_rule_t, batch_rule_t batch_rule>
  13360. at::Tensor _to_sparse_csc_generated_plumbing(const at::Tensor & self, ::std::optional<int64_t> dense_dim) {
  13361. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13362. auto maybe_layer = maybeCurrentDynamicLayer();
  13363. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13364. int64_t cur_level = maybe_layer->layerId();
  13365. if (!isBatchedAtLevel(self, cur_level)) {
  13366. return at::_ops::_to_sparse_csc::call(self, dense_dim);
  13367. }
  13368. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13369. auto results = batch_rule(self_value, self_bdim, dense_dim);
  13370. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13371. }
  13372. template <typename batch_rule_t, batch_rule_t batch_rule>
  13373. at::Tensor to_sparse_bsr_generated_plumbing(const at::Tensor & self, at::IntArrayRef blocksize, ::std::optional<int64_t> dense_dim) {
  13374. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13375. auto maybe_layer = maybeCurrentDynamicLayer();
  13376. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13377. int64_t cur_level = maybe_layer->layerId();
  13378. if (!isBatchedAtLevel(self, cur_level)) {
  13379. return at::_ops::to_sparse_bsr::call(self, blocksize, dense_dim);
  13380. }
  13381. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13382. auto results = batch_rule(self_value, self_bdim, blocksize, dense_dim);
  13383. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13384. }
  13385. template <typename batch_rule_t, batch_rule_t batch_rule>
  13386. at::Tensor _to_sparse_bsr_generated_plumbing(const at::Tensor & self, at::IntArrayRef blocksize, ::std::optional<int64_t> dense_dim) {
  13387. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13388. auto maybe_layer = maybeCurrentDynamicLayer();
  13389. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13390. int64_t cur_level = maybe_layer->layerId();
  13391. if (!isBatchedAtLevel(self, cur_level)) {
  13392. return at::_ops::_to_sparse_bsr::call(self, blocksize, dense_dim);
  13393. }
  13394. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13395. auto results = batch_rule(self_value, self_bdim, blocksize, dense_dim);
  13396. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13397. }
  13398. template <typename batch_rule_t, batch_rule_t batch_rule>
  13399. at::Tensor to_sparse_bsc_generated_plumbing(const at::Tensor & self, at::IntArrayRef blocksize, ::std::optional<int64_t> dense_dim) {
  13400. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13401. auto maybe_layer = maybeCurrentDynamicLayer();
  13402. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13403. int64_t cur_level = maybe_layer->layerId();
  13404. if (!isBatchedAtLevel(self, cur_level)) {
  13405. return at::_ops::to_sparse_bsc::call(self, blocksize, dense_dim);
  13406. }
  13407. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13408. auto results = batch_rule(self_value, self_bdim, blocksize, dense_dim);
  13409. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13410. }
  13411. template <typename batch_rule_t, batch_rule_t batch_rule>
  13412. at::Tensor _to_sparse_bsc_generated_plumbing(const at::Tensor & self, at::IntArrayRef blocksize, ::std::optional<int64_t> dense_dim) {
  13413. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13414. auto maybe_layer = maybeCurrentDynamicLayer();
  13415. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13416. int64_t cur_level = maybe_layer->layerId();
  13417. if (!isBatchedAtLevel(self, cur_level)) {
  13418. return at::_ops::_to_sparse_bsc::call(self, blocksize, dense_dim);
  13419. }
  13420. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13421. auto results = batch_rule(self_value, self_bdim, blocksize, dense_dim);
  13422. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13423. }
  13424. template <typename batch_rule_t, batch_rule_t batch_rule>
  13425. ::std::tuple<at::Tensor,at::Tensor> _to_sparse_semi_structured_generated_plumbing(const at::Tensor & dense) {
  13426. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13427. auto maybe_layer = maybeCurrentDynamicLayer();
  13428. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13429. int64_t cur_level = maybe_layer->layerId();
  13430. if (!isBatchedAtLevel(dense, cur_level)) {
  13431. return at::_ops::_to_sparse_semi_structured::call(dense);
  13432. }
  13433. auto [dense_value, dense_bdim] = unwrapTensorAtLevel(dense, cur_level);
  13434. auto results = batch_rule(dense_value, dense_bdim);
  13435. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  13436. }
  13437. template <typename batch_rule_t, batch_rule_t batch_rule>
  13438. at::Tensor to_mkldnn_generated_plumbing(const at::Tensor & self, ::std::optional<at::ScalarType> dtype) {
  13439. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13440. auto maybe_layer = maybeCurrentDynamicLayer();
  13441. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13442. int64_t cur_level = maybe_layer->layerId();
  13443. if (!isBatchedAtLevel(self, cur_level)) {
  13444. return at::_ops::to_mkldnn::call(self, dtype);
  13445. }
  13446. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13447. auto results = batch_rule(self_value, self_bdim, dtype);
  13448. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13449. }
  13450. template <typename batch_rule_t, batch_rule_t batch_rule>
  13451. at::Tensor mkldnn_reorder_conv2d_weight_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef padding, c10::SymIntArrayRef stride, c10::SymIntArrayRef dilation, c10::SymInt groups, at::OptionalSymIntArrayRef input_size) {
  13452. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13453. auto maybe_layer = maybeCurrentDynamicLayer();
  13454. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13455. int64_t cur_level = maybe_layer->layerId();
  13456. if (!isBatchedAtLevel(self, cur_level)) {
  13457. return at::_ops::mkldnn_reorder_conv2d_weight::call(self, padding, stride, dilation, groups, input_size);
  13458. }
  13459. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13460. auto results = batch_rule(self_value, self_bdim, padding, stride, dilation, groups, input_size);
  13461. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13462. }
  13463. template <typename batch_rule_t, batch_rule_t batch_rule>
  13464. at::Tensor mkldnn_reorder_conv3d_weight_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef padding, c10::SymIntArrayRef stride, c10::SymIntArrayRef dilation, c10::SymInt groups, at::OptionalSymIntArrayRef input_size) {
  13465. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13466. auto maybe_layer = maybeCurrentDynamicLayer();
  13467. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13468. int64_t cur_level = maybe_layer->layerId();
  13469. if (!isBatchedAtLevel(self, cur_level)) {
  13470. return at::_ops::mkldnn_reorder_conv3d_weight::call(self, padding, stride, dilation, groups, input_size);
  13471. }
  13472. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13473. auto results = batch_rule(self_value, self_bdim, padding, stride, dilation, groups, input_size);
  13474. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13475. }
  13476. template <typename batch_rule_t, batch_rule_t batch_rule>
  13477. at::Tensor to_mkldnn_backward_generated_plumbing(const at::Tensor & grad, const at::Tensor & input) {
  13478. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13479. auto maybe_layer = maybeCurrentDynamicLayer();
  13480. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13481. int64_t cur_level = maybe_layer->layerId();
  13482. if (!isBatchedAtLevel(grad, cur_level) && !isBatchedAtLevel(input, cur_level)) {
  13483. return at::_ops::to_mkldnn_backward::call(grad, input);
  13484. }
  13485. auto [grad_value, grad_bdim] = unwrapTensorAtLevel(grad, cur_level);
  13486. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  13487. auto results = batch_rule(grad_value, grad_bdim, input_value, input_bdim);
  13488. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13489. }
  13490. template <typename batch_rule_t, batch_rule_t batch_rule>
  13491. at::Tensor quantize_per_tensor_dynamic_generated_plumbing(const at::Tensor & self, at::ScalarType dtype, bool reduce_range) {
  13492. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13493. auto maybe_layer = maybeCurrentDynamicLayer();
  13494. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13495. int64_t cur_level = maybe_layer->layerId();
  13496. if (!isBatchedAtLevel(self, cur_level)) {
  13497. return at::_ops::quantize_per_tensor_dynamic::call(self, dtype, reduce_range);
  13498. }
  13499. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13500. auto results = batch_rule(self_value, self_bdim, dtype, reduce_range);
  13501. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13502. }
  13503. template <typename batch_rule_t, batch_rule_t batch_rule>
  13504. at::Tensor quantize_per_tensor_generated_plumbing(const at::Tensor & self, double scale, int64_t zero_point, at::ScalarType dtype) {
  13505. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13506. auto maybe_layer = maybeCurrentDynamicLayer();
  13507. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13508. int64_t cur_level = maybe_layer->layerId();
  13509. if (!isBatchedAtLevel(self, cur_level)) {
  13510. return at::_ops::quantize_per_tensor::call(self, scale, zero_point, dtype);
  13511. }
  13512. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13513. auto results = batch_rule(self_value, self_bdim, scale, zero_point, dtype);
  13514. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13515. }
  13516. template <typename batch_rule_t, batch_rule_t batch_rule>
  13517. at::Tensor quantize_per_tensor_tensor_qparams_generated_plumbing(const at::Tensor & self, const at::Tensor & scale, const at::Tensor & zero_point, at::ScalarType dtype) {
  13518. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13519. auto maybe_layer = maybeCurrentDynamicLayer();
  13520. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13521. int64_t cur_level = maybe_layer->layerId();
  13522. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(scale, cur_level) && !isBatchedAtLevel(zero_point, cur_level)) {
  13523. return at::_ops::quantize_per_tensor_tensor_qparams::call(self, scale, zero_point, dtype);
  13524. }
  13525. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13526. auto [scale_value, scale_bdim] = unwrapTensorAtLevel(scale, cur_level);
  13527. auto [zero_point_value, zero_point_bdim] = unwrapTensorAtLevel(zero_point, cur_level);
  13528. auto results = batch_rule(self_value, self_bdim, scale_value, scale_bdim, zero_point_value, zero_point_bdim, dtype);
  13529. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13530. }
  13531. template <typename batch_rule_t, batch_rule_t batch_rule>
  13532. ::std::vector<at::Tensor> quantize_per_tensor_tensors_generated_plumbing(at::TensorList tensors, const at::Tensor & scales, const at::Tensor & zero_points, at::ScalarType dtype) {
  13533. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13534. auto maybe_layer = maybeCurrentDynamicLayer();
  13535. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13536. int64_t cur_level = maybe_layer->layerId();
  13537. if (!isBatchedAtLevel(tensors, cur_level) && !isBatchedAtLevel(scales, cur_level) && !isBatchedAtLevel(zero_points, cur_level)) {
  13538. return at::_ops::quantize_per_tensor_tensors::call(tensors, scales, zero_points, dtype);
  13539. }
  13540. auto [scales_value, scales_bdim] = unwrapTensorAtLevel(scales, cur_level);
  13541. auto [zero_points_value, zero_points_bdim] = unwrapTensorAtLevel(zero_points, cur_level);
  13542. auto results = batch_rule(tensors, scales_value, scales_bdim, zero_points_value, zero_points_bdim, dtype);
  13543. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  13544. }
  13545. template <typename batch_rule_t, batch_rule_t batch_rule>
  13546. at::Tensor quantize_per_channel_generated_plumbing(const at::Tensor & self, const at::Tensor & scales, const at::Tensor & zero_points, int64_t axis, at::ScalarType dtype) {
  13547. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13548. auto maybe_layer = maybeCurrentDynamicLayer();
  13549. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13550. int64_t cur_level = maybe_layer->layerId();
  13551. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(scales, cur_level) && !isBatchedAtLevel(zero_points, cur_level)) {
  13552. return at::_ops::quantize_per_channel::call(self, scales, zero_points, axis, dtype);
  13553. }
  13554. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13555. auto [scales_value, scales_bdim] = unwrapTensorAtLevel(scales, cur_level);
  13556. auto [zero_points_value, zero_points_bdim] = unwrapTensorAtLevel(zero_points, cur_level);
  13557. auto results = batch_rule(self_value, self_bdim, scales_value, scales_bdim, zero_points_value, zero_points_bdim, axis, dtype);
  13558. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13559. }
  13560. template <typename batch_rule_t, batch_rule_t batch_rule>
  13561. at::Tensor dequantize_self_generated_plumbing(const at::Tensor & self) {
  13562. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13563. auto maybe_layer = maybeCurrentDynamicLayer();
  13564. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13565. int64_t cur_level = maybe_layer->layerId();
  13566. if (!isBatchedAtLevel(self, cur_level)) {
  13567. return at::_ops::dequantize_self::call(self);
  13568. }
  13569. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13570. auto results = batch_rule(self_value, self_bdim);
  13571. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13572. }
  13573. template <typename batch_rule_t, batch_rule_t batch_rule>
  13574. ::std::vector<at::Tensor> dequantize_tensors_generated_plumbing(at::TensorList tensors) {
  13575. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13576. auto maybe_layer = maybeCurrentDynamicLayer();
  13577. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13578. int64_t cur_level = maybe_layer->layerId();
  13579. if (!isBatchedAtLevel(tensors, cur_level)) {
  13580. return at::_ops::dequantize_tensors::call(tensors);
  13581. }
  13582. auto results = batch_rule(tensors);
  13583. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  13584. }
  13585. template <typename batch_rule_t, batch_rule_t batch_rule>
  13586. at::Tensor q_per_channel_scales_generated_plumbing(const at::Tensor & self) {
  13587. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13588. auto maybe_layer = maybeCurrentDynamicLayer();
  13589. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13590. int64_t cur_level = maybe_layer->layerId();
  13591. if (!isBatchedAtLevel(self, cur_level)) {
  13592. return at::_ops::q_per_channel_scales::call(self);
  13593. }
  13594. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13595. auto results = batch_rule(self_value, self_bdim);
  13596. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13597. }
  13598. template <typename batch_rule_t, batch_rule_t batch_rule>
  13599. at::Tensor q_per_channel_zero_points_generated_plumbing(const at::Tensor & self) {
  13600. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13601. auto maybe_layer = maybeCurrentDynamicLayer();
  13602. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13603. int64_t cur_level = maybe_layer->layerId();
  13604. if (!isBatchedAtLevel(self, cur_level)) {
  13605. return at::_ops::q_per_channel_zero_points::call(self);
  13606. }
  13607. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13608. auto results = batch_rule(self_value, self_bdim);
  13609. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13610. }
  13611. template <typename batch_rule_t, batch_rule_t batch_rule>
  13612. at::Tensor int_repr_generated_plumbing(const at::Tensor & self) {
  13613. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13614. auto maybe_layer = maybeCurrentDynamicLayer();
  13615. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13616. int64_t cur_level = maybe_layer->layerId();
  13617. if (!isBatchedAtLevel(self, cur_level)) {
  13618. return at::_ops::int_repr::call(self);
  13619. }
  13620. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13621. auto results = batch_rule(self_value, self_bdim);
  13622. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13623. }
  13624. template <typename batch_rule_t, batch_rule_t batch_rule>
  13625. at::Tensor _make_per_tensor_quantized_tensor_generated_plumbing(const at::Tensor & self, double scale, int64_t zero_point) {
  13626. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13627. auto maybe_layer = maybeCurrentDynamicLayer();
  13628. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13629. int64_t cur_level = maybe_layer->layerId();
  13630. if (!isBatchedAtLevel(self, cur_level)) {
  13631. return at::_ops::_make_per_tensor_quantized_tensor::call(self, scale, zero_point);
  13632. }
  13633. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13634. auto results = batch_rule(self_value, self_bdim, scale, zero_point);
  13635. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13636. }
  13637. template <typename batch_rule_t, batch_rule_t batch_rule>
  13638. at::Tensor _make_per_channel_quantized_tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & scale, const at::Tensor & zero_point, int64_t axis) {
  13639. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13640. auto maybe_layer = maybeCurrentDynamicLayer();
  13641. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13642. int64_t cur_level = maybe_layer->layerId();
  13643. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(scale, cur_level) && !isBatchedAtLevel(zero_point, cur_level)) {
  13644. return at::_ops::_make_per_channel_quantized_tensor::call(self, scale, zero_point, axis);
  13645. }
  13646. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13647. auto [scale_value, scale_bdim] = unwrapTensorAtLevel(scale, cur_level);
  13648. auto [zero_point_value, zero_point_bdim] = unwrapTensorAtLevel(zero_point, cur_level);
  13649. auto results = batch_rule(self_value, self_bdim, scale_value, scale_bdim, zero_point_value, zero_point_bdim, axis);
  13650. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13651. }
  13652. template <typename batch_rule_t, batch_rule_t batch_rule>
  13653. at::Tensor fake_quantize_per_tensor_affine_generated_plumbing(const at::Tensor & self, double scale, int64_t zero_point, int64_t quant_min, int64_t quant_max) {
  13654. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13655. auto maybe_layer = maybeCurrentDynamicLayer();
  13656. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13657. int64_t cur_level = maybe_layer->layerId();
  13658. if (!isBatchedAtLevel(self, cur_level)) {
  13659. return at::_ops::fake_quantize_per_tensor_affine::call(self, scale, zero_point, quant_min, quant_max);
  13660. }
  13661. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13662. auto results = batch_rule(self_value, self_bdim, scale, zero_point, quant_min, quant_max);
  13663. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13664. }
  13665. template <typename batch_rule_t, batch_rule_t batch_rule>
  13666. at::Tensor fake_quantize_per_tensor_affine_tensor_qparams_generated_plumbing(const at::Tensor & self, const at::Tensor & scale, const at::Tensor & zero_point, int64_t quant_min, int64_t quant_max) {
  13667. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13668. auto maybe_layer = maybeCurrentDynamicLayer();
  13669. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13670. int64_t cur_level = maybe_layer->layerId();
  13671. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(scale, cur_level) && !isBatchedAtLevel(zero_point, cur_level)) {
  13672. return at::_ops::fake_quantize_per_tensor_affine_tensor_qparams::call(self, scale, zero_point, quant_min, quant_max);
  13673. }
  13674. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13675. auto [scale_value, scale_bdim] = unwrapTensorAtLevel(scale, cur_level);
  13676. auto [zero_point_value, zero_point_bdim] = unwrapTensorAtLevel(zero_point, cur_level);
  13677. auto results = batch_rule(self_value, self_bdim, scale_value, scale_bdim, zero_point_value, zero_point_bdim, quant_min, quant_max);
  13678. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13679. }
  13680. template <typename batch_rule_t, batch_rule_t batch_rule>
  13681. ::std::tuple<at::Tensor,at::Tensor> fake_quantize_per_tensor_affine_cachemask_generated_plumbing(const at::Tensor & self, double scale, int64_t zero_point, int64_t quant_min, int64_t quant_max) {
  13682. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13683. auto maybe_layer = maybeCurrentDynamicLayer();
  13684. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13685. int64_t cur_level = maybe_layer->layerId();
  13686. if (!isBatchedAtLevel(self, cur_level)) {
  13687. return at::_ops::fake_quantize_per_tensor_affine_cachemask::call(self, scale, zero_point, quant_min, quant_max);
  13688. }
  13689. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13690. auto results = batch_rule(self_value, self_bdim, scale, zero_point, quant_min, quant_max);
  13691. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  13692. }
  13693. template <typename batch_rule_t, batch_rule_t batch_rule>
  13694. ::std::tuple<at::Tensor,at::Tensor> _fake_quantize_per_tensor_affine_cachemask_tensor_qparams_generated_plumbing(const at::Tensor & self, const at::Tensor & scale, const at::Tensor & zero_point, const at::Tensor & fake_quant_enabled, int64_t quant_min, int64_t quant_max) {
  13695. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13696. auto maybe_layer = maybeCurrentDynamicLayer();
  13697. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13698. int64_t cur_level = maybe_layer->layerId();
  13699. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(scale, cur_level) && !isBatchedAtLevel(zero_point, cur_level) && !isBatchedAtLevel(fake_quant_enabled, cur_level)) {
  13700. return at::_ops::_fake_quantize_per_tensor_affine_cachemask_tensor_qparams::call(self, scale, zero_point, fake_quant_enabled, quant_min, quant_max);
  13701. }
  13702. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13703. auto [scale_value, scale_bdim] = unwrapTensorAtLevel(scale, cur_level);
  13704. auto [zero_point_value, zero_point_bdim] = unwrapTensorAtLevel(zero_point, cur_level);
  13705. auto [fake_quant_enabled_value, fake_quant_enabled_bdim] = unwrapTensorAtLevel(fake_quant_enabled, cur_level);
  13706. auto results = batch_rule(self_value, self_bdim, scale_value, scale_bdim, zero_point_value, zero_point_bdim, fake_quant_enabled_value, fake_quant_enabled_bdim, quant_min, quant_max);
  13707. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  13708. }
  13709. template <typename batch_rule_t, batch_rule_t batch_rule>
  13710. at::Tensor fake_quantize_per_tensor_affine_cachemask_backward_generated_plumbing(const at::Tensor & grad, const at::Tensor & mask) {
  13711. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13712. auto maybe_layer = maybeCurrentDynamicLayer();
  13713. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13714. int64_t cur_level = maybe_layer->layerId();
  13715. if (!isBatchedAtLevel(grad, cur_level) && !isBatchedAtLevel(mask, cur_level)) {
  13716. return at::_ops::fake_quantize_per_tensor_affine_cachemask_backward::call(grad, mask);
  13717. }
  13718. auto [grad_value, grad_bdim] = unwrapTensorAtLevel(grad, cur_level);
  13719. auto [mask_value, mask_bdim] = unwrapTensorAtLevel(mask, cur_level);
  13720. auto results = batch_rule(grad_value, grad_bdim, mask_value, mask_bdim);
  13721. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13722. }
  13723. template <typename batch_rule_t, batch_rule_t batch_rule>
  13724. at::Tensor _fake_quantize_learnable_per_tensor_affine_generated_plumbing(const at::Tensor & self, const at::Tensor & scale, const at::Tensor & zero_point, int64_t quant_min, int64_t quant_max, double grad_factor) {
  13725. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13726. auto maybe_layer = maybeCurrentDynamicLayer();
  13727. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13728. int64_t cur_level = maybe_layer->layerId();
  13729. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(scale, cur_level) && !isBatchedAtLevel(zero_point, cur_level)) {
  13730. return at::_ops::_fake_quantize_learnable_per_tensor_affine::call(self, scale, zero_point, quant_min, quant_max, grad_factor);
  13731. }
  13732. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13733. auto [scale_value, scale_bdim] = unwrapTensorAtLevel(scale, cur_level);
  13734. auto [zero_point_value, zero_point_bdim] = unwrapTensorAtLevel(zero_point, cur_level);
  13735. auto results = batch_rule(self_value, self_bdim, scale_value, scale_bdim, zero_point_value, zero_point_bdim, quant_min, quant_max, grad_factor);
  13736. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13737. }
  13738. template <typename batch_rule_t, batch_rule_t batch_rule>
  13739. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> _fake_quantize_learnable_per_tensor_affine_backward_generated_plumbing(const at::Tensor & grad, const at::Tensor & self, const at::Tensor & scale, const at::Tensor & zero_point, int64_t quant_min, int64_t quant_max, double grad_factor) {
  13740. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13741. auto maybe_layer = maybeCurrentDynamicLayer();
  13742. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13743. int64_t cur_level = maybe_layer->layerId();
  13744. if (!isBatchedAtLevel(grad, cur_level) && !isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(scale, cur_level) && !isBatchedAtLevel(zero_point, cur_level)) {
  13745. return at::_ops::_fake_quantize_learnable_per_tensor_affine_backward::call(grad, self, scale, zero_point, quant_min, quant_max, grad_factor);
  13746. }
  13747. auto [grad_value, grad_bdim] = unwrapTensorAtLevel(grad, cur_level);
  13748. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13749. auto [scale_value, scale_bdim] = unwrapTensorAtLevel(scale, cur_level);
  13750. auto [zero_point_value, zero_point_bdim] = unwrapTensorAtLevel(zero_point, cur_level);
  13751. auto results = batch_rule(grad_value, grad_bdim, self_value, self_bdim, scale_value, scale_bdim, zero_point_value, zero_point_bdim, quant_min, quant_max, grad_factor);
  13752. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  13753. }
  13754. template <typename batch_rule_t, batch_rule_t batch_rule>
  13755. at::Tensor fake_quantize_per_channel_affine_generated_plumbing(const at::Tensor & self, const at::Tensor & scale, const at::Tensor & zero_point, int64_t axis, int64_t quant_min, int64_t quant_max) {
  13756. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13757. auto maybe_layer = maybeCurrentDynamicLayer();
  13758. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13759. int64_t cur_level = maybe_layer->layerId();
  13760. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(scale, cur_level) && !isBatchedAtLevel(zero_point, cur_level)) {
  13761. return at::_ops::fake_quantize_per_channel_affine::call(self, scale, zero_point, axis, quant_min, quant_max);
  13762. }
  13763. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13764. auto [scale_value, scale_bdim] = unwrapTensorAtLevel(scale, cur_level);
  13765. auto [zero_point_value, zero_point_bdim] = unwrapTensorAtLevel(zero_point, cur_level);
  13766. auto results = batch_rule(self_value, self_bdim, scale_value, scale_bdim, zero_point_value, zero_point_bdim, axis, quant_min, quant_max);
  13767. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13768. }
  13769. template <typename batch_rule_t, batch_rule_t batch_rule>
  13770. ::std::tuple<at::Tensor,at::Tensor> fake_quantize_per_channel_affine_cachemask_generated_plumbing(const at::Tensor & self, const at::Tensor & scale, const at::Tensor & zero_point, int64_t axis, int64_t quant_min, int64_t quant_max) {
  13771. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13772. auto maybe_layer = maybeCurrentDynamicLayer();
  13773. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13774. int64_t cur_level = maybe_layer->layerId();
  13775. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(scale, cur_level) && !isBatchedAtLevel(zero_point, cur_level)) {
  13776. return at::_ops::fake_quantize_per_channel_affine_cachemask::call(self, scale, zero_point, axis, quant_min, quant_max);
  13777. }
  13778. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13779. auto [scale_value, scale_bdim] = unwrapTensorAtLevel(scale, cur_level);
  13780. auto [zero_point_value, zero_point_bdim] = unwrapTensorAtLevel(zero_point, cur_level);
  13781. auto results = batch_rule(self_value, self_bdim, scale_value, scale_bdim, zero_point_value, zero_point_bdim, axis, quant_min, quant_max);
  13782. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  13783. }
  13784. template <typename batch_rule_t, batch_rule_t batch_rule>
  13785. at::Tensor fake_quantize_per_channel_affine_cachemask_backward_generated_plumbing(const at::Tensor & grad, const at::Tensor & mask) {
  13786. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13787. auto maybe_layer = maybeCurrentDynamicLayer();
  13788. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13789. int64_t cur_level = maybe_layer->layerId();
  13790. if (!isBatchedAtLevel(grad, cur_level) && !isBatchedAtLevel(mask, cur_level)) {
  13791. return at::_ops::fake_quantize_per_channel_affine_cachemask_backward::call(grad, mask);
  13792. }
  13793. auto [grad_value, grad_bdim] = unwrapTensorAtLevel(grad, cur_level);
  13794. auto [mask_value, mask_bdim] = unwrapTensorAtLevel(mask, cur_level);
  13795. auto results = batch_rule(grad_value, grad_bdim, mask_value, mask_bdim);
  13796. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13797. }
  13798. template <typename batch_rule_t, batch_rule_t batch_rule>
  13799. at::Tensor _fake_quantize_learnable_per_channel_affine_generated_plumbing(const at::Tensor & self, const at::Tensor & scale, const at::Tensor & zero_point, int64_t axis, int64_t quant_min, int64_t quant_max, double grad_factor) {
  13800. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13801. auto maybe_layer = maybeCurrentDynamicLayer();
  13802. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13803. int64_t cur_level = maybe_layer->layerId();
  13804. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(scale, cur_level) && !isBatchedAtLevel(zero_point, cur_level)) {
  13805. return at::_ops::_fake_quantize_learnable_per_channel_affine::call(self, scale, zero_point, axis, quant_min, quant_max, grad_factor);
  13806. }
  13807. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13808. auto [scale_value, scale_bdim] = unwrapTensorAtLevel(scale, cur_level);
  13809. auto [zero_point_value, zero_point_bdim] = unwrapTensorAtLevel(zero_point, cur_level);
  13810. auto results = batch_rule(self_value, self_bdim, scale_value, scale_bdim, zero_point_value, zero_point_bdim, axis, quant_min, quant_max, grad_factor);
  13811. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13812. }
  13813. template <typename batch_rule_t, batch_rule_t batch_rule>
  13814. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> _fake_quantize_learnable_per_channel_affine_backward_generated_plumbing(const at::Tensor & grad, const at::Tensor & self, const at::Tensor & scale, const at::Tensor & zero_point, int64_t axis, int64_t quant_min, int64_t quant_max, double grad_factor) {
  13815. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13816. auto maybe_layer = maybeCurrentDynamicLayer();
  13817. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13818. int64_t cur_level = maybe_layer->layerId();
  13819. if (!isBatchedAtLevel(grad, cur_level) && !isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(scale, cur_level) && !isBatchedAtLevel(zero_point, cur_level)) {
  13820. return at::_ops::_fake_quantize_learnable_per_channel_affine_backward::call(grad, self, scale, zero_point, axis, quant_min, quant_max, grad_factor);
  13821. }
  13822. auto [grad_value, grad_bdim] = unwrapTensorAtLevel(grad, cur_level);
  13823. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13824. auto [scale_value, scale_bdim] = unwrapTensorAtLevel(scale, cur_level);
  13825. auto [zero_point_value, zero_point_bdim] = unwrapTensorAtLevel(zero_point, cur_level);
  13826. auto results = batch_rule(grad_value, grad_bdim, self_value, self_bdim, scale_value, scale_bdim, zero_point_value, zero_point_bdim, axis, quant_min, quant_max, grad_factor);
  13827. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  13828. }
  13829. template <typename batch_rule_t, batch_rule_t batch_rule>
  13830. at::Tensor _saturate_weight_to_fp16_generated_plumbing(const at::Tensor & weight) {
  13831. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13832. auto maybe_layer = maybeCurrentDynamicLayer();
  13833. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13834. int64_t cur_level = maybe_layer->layerId();
  13835. if (!isBatchedAtLevel(weight, cur_level)) {
  13836. return at::_ops::_saturate_weight_to_fp16::call(weight);
  13837. }
  13838. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  13839. auto results = batch_rule(weight_value, weight_bdim);
  13840. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13841. }
  13842. template <typename batch_rule_t, batch_rule_t batch_rule>
  13843. ::std::tuple<at::Tensor,at::Tensor> choose_qparams_optimized_generated_plumbing(const at::Tensor & input, int64_t numel, int64_t n_bins, double ratio, int64_t bit_width) {
  13844. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13845. auto maybe_layer = maybeCurrentDynamicLayer();
  13846. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13847. int64_t cur_level = maybe_layer->layerId();
  13848. if (!isBatchedAtLevel(input, cur_level)) {
  13849. return at::_ops::choose_qparams_optimized::call(input, numel, n_bins, ratio, bit_width);
  13850. }
  13851. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  13852. auto results = batch_rule(input_value, input_bdim, numel, n_bins, ratio, bit_width);
  13853. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  13854. }
  13855. template <typename batch_rule_t, batch_rule_t batch_rule>
  13856. at::Tensor _autocast_to_reduced_precision_generated_plumbing(const at::Tensor & self, bool cuda_enabled, bool cpu_enabled, at::ScalarType cuda_dtype, at::ScalarType cpu_dtype) {
  13857. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13858. auto maybe_layer = maybeCurrentDynamicLayer();
  13859. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13860. int64_t cur_level = maybe_layer->layerId();
  13861. if (!isBatchedAtLevel(self, cur_level)) {
  13862. return at::_ops::_autocast_to_reduced_precision::call(self, cuda_enabled, cpu_enabled, cuda_dtype, cpu_dtype);
  13863. }
  13864. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13865. auto results = batch_rule(self_value, self_bdim, cuda_enabled, cpu_enabled, cuda_dtype, cpu_dtype);
  13866. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13867. }
  13868. template <typename batch_rule_t, batch_rule_t batch_rule>
  13869. at::Tensor _autocast_to_full_precision_generated_plumbing(const at::Tensor & self, bool cuda_enabled, bool cpu_enabled) {
  13870. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13871. auto maybe_layer = maybeCurrentDynamicLayer();
  13872. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13873. int64_t cur_level = maybe_layer->layerId();
  13874. if (!isBatchedAtLevel(self, cur_level)) {
  13875. return at::_ops::_autocast_to_full_precision::call(self, cuda_enabled, cpu_enabled);
  13876. }
  13877. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13878. auto results = batch_rule(self_value, self_bdim, cuda_enabled, cpu_enabled);
  13879. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13880. }
  13881. template <typename batch_rule_t, batch_rule_t batch_rule>
  13882. at::Tensor _to_copy_generated_plumbing(const at::Tensor & self, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory, bool non_blocking, ::std::optional<at::MemoryFormat> memory_format) {
  13883. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13884. auto maybe_layer = maybeCurrentDynamicLayer();
  13885. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13886. int64_t cur_level = maybe_layer->layerId();
  13887. if (!isBatchedAtLevel(self, cur_level)) {
  13888. return at::_ops::_to_copy::call(self, dtype, layout, device, pin_memory, non_blocking, memory_format);
  13889. }
  13890. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13891. auto results = batch_rule(self_value, self_bdim, dtype, layout, device, pin_memory, non_blocking, memory_format);
  13892. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13893. }
  13894. template <typename batch_rule_t, batch_rule_t batch_rule>
  13895. at::Tensor to_dtype_layout_generated_plumbing(const at::Tensor & self, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory, bool non_blocking, bool copy, ::std::optional<at::MemoryFormat> memory_format) {
  13896. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13897. auto maybe_layer = maybeCurrentDynamicLayer();
  13898. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13899. int64_t cur_level = maybe_layer->layerId();
  13900. if (!isBatchedAtLevel(self, cur_level)) {
  13901. return at::_ops::to_dtype_layout::call(self, dtype, layout, device, pin_memory, non_blocking, copy, memory_format);
  13902. }
  13903. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13904. auto results = batch_rule(self_value, self_bdim, dtype, layout, device, pin_memory, non_blocking, copy, memory_format);
  13905. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13906. }
  13907. template <typename batch_rule_t, batch_rule_t batch_rule>
  13908. at::Tensor to_device_generated_plumbing(const at::Tensor & self, at::Device device, at::ScalarType dtype, bool non_blocking, bool copy, ::std::optional<at::MemoryFormat> memory_format) {
  13909. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13910. auto maybe_layer = maybeCurrentDynamicLayer();
  13911. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13912. int64_t cur_level = maybe_layer->layerId();
  13913. if (!isBatchedAtLevel(self, cur_level)) {
  13914. return at::_ops::to_device::call(self, device, dtype, non_blocking, copy, memory_format);
  13915. }
  13916. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13917. auto results = batch_rule(self_value, self_bdim, device, dtype, non_blocking, copy, memory_format);
  13918. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13919. }
  13920. template <typename batch_rule_t, batch_rule_t batch_rule>
  13921. at::Tensor to_dtype_generated_plumbing(const at::Tensor & self, at::ScalarType dtype, bool non_blocking, bool copy, ::std::optional<at::MemoryFormat> memory_format) {
  13922. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13923. auto maybe_layer = maybeCurrentDynamicLayer();
  13924. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13925. int64_t cur_level = maybe_layer->layerId();
  13926. if (!isBatchedAtLevel(self, cur_level)) {
  13927. return at::_ops::to_dtype::call(self, dtype, non_blocking, copy, memory_format);
  13928. }
  13929. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13930. auto results = batch_rule(self_value, self_bdim, dtype, non_blocking, copy, memory_format);
  13931. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13932. }
  13933. template <typename batch_rule_t, batch_rule_t batch_rule>
  13934. at::Tensor to_other_generated_plumbing(const at::Tensor & self, const at::Tensor & other, bool non_blocking, bool copy, ::std::optional<at::MemoryFormat> memory_format) {
  13935. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13936. auto maybe_layer = maybeCurrentDynamicLayer();
  13937. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13938. int64_t cur_level = maybe_layer->layerId();
  13939. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  13940. return at::_ops::to_other::call(self, other, non_blocking, copy, memory_format);
  13941. }
  13942. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13943. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  13944. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim, non_blocking, copy, memory_format);
  13945. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13946. }
  13947. template <typename batch_rule_t, batch_rule_t batch_rule>
  13948. ::std::vector<at::Tensor> meshgrid_generated_plumbing(at::TensorList tensors) {
  13949. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13950. auto maybe_layer = maybeCurrentDynamicLayer();
  13951. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13952. int64_t cur_level = maybe_layer->layerId();
  13953. if (!isBatchedAtLevel(tensors, cur_level)) {
  13954. return at::_ops::meshgrid::call(tensors);
  13955. }
  13956. auto results = batch_rule(tensors);
  13957. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  13958. }
  13959. template <typename batch_rule_t, batch_rule_t batch_rule>
  13960. ::std::vector<at::Tensor> meshgrid_indexing_generated_plumbing(at::TensorList tensors, c10::string_view indexing) {
  13961. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13962. auto maybe_layer = maybeCurrentDynamicLayer();
  13963. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13964. int64_t cur_level = maybe_layer->layerId();
  13965. if (!isBatchedAtLevel(tensors, cur_level)) {
  13966. return at::_ops::meshgrid_indexing::call(tensors, indexing);
  13967. }
  13968. auto results = batch_rule(tensors, indexing);
  13969. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  13970. }
  13971. template <typename batch_rule_t, batch_rule_t batch_rule>
  13972. at::Tensor cartesian_prod_generated_plumbing(at::TensorList tensors) {
  13973. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13974. auto maybe_layer = maybeCurrentDynamicLayer();
  13975. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13976. int64_t cur_level = maybe_layer->layerId();
  13977. if (!isBatchedAtLevel(tensors, cur_level)) {
  13978. return at::_ops::cartesian_prod::call(tensors);
  13979. }
  13980. auto results = batch_rule(tensors);
  13981. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13982. }
  13983. template <typename batch_rule_t, batch_rule_t batch_rule>
  13984. at::Tensor combinations_generated_plumbing(const at::Tensor & self, int64_t r, bool with_replacement) {
  13985. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13986. auto maybe_layer = maybeCurrentDynamicLayer();
  13987. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  13988. int64_t cur_level = maybe_layer->layerId();
  13989. if (!isBatchedAtLevel(self, cur_level)) {
  13990. return at::_ops::combinations::call(self, r, with_replacement);
  13991. }
  13992. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  13993. auto results = batch_rule(self_value, self_bdim, r, with_replacement);
  13994. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  13995. }
  13996. template <typename batch_rule_t, batch_rule_t batch_rule>
  13997. ::std::tuple<at::Tensor,at::Tensor,at::Tensor,at::Tensor,at::Tensor,at::Tensor> _lstm_mps_generated_plumbing(const at::Tensor & input, at::TensorList hx, at::TensorList params, bool has_biases, int64_t num_layers, double dropout, bool train, bool bidirectional, bool batch_first) {
  13998. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  13999. auto maybe_layer = maybeCurrentDynamicLayer();
  14000. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14001. int64_t cur_level = maybe_layer->layerId();
  14002. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(hx, cur_level) && !isBatchedAtLevel(params, cur_level)) {
  14003. return at::_ops::_lstm_mps::call(input, hx, params, has_biases, num_layers, dropout, train, bidirectional, batch_first);
  14004. }
  14005. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  14006. auto results = batch_rule(input_value, input_bdim, hx, params, has_biases, num_layers, dropout, train, bidirectional, batch_first);
  14007. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level), makeBatched(std::get<6>(results), std::get<7>(results), cur_level), makeBatched(std::get<8>(results), std::get<9>(results), cur_level), makeBatched(std::get<10>(results), std::get<11>(results), cur_level));
  14008. }
  14009. template <typename batch_rule_t, batch_rule_t batch_rule>
  14010. ::std::tuple<at::Tensor,::std::vector<at::Tensor>,::std::vector<at::Tensor>> lstm_mps_backward_generated_plumbing(const ::std::optional<at::Tensor> & grad_y, const ::std::optional<at::Tensor> & grad_hy, const ::std::optional<at::Tensor> & grad_cy, const at::Tensor & z_state, const at::Tensor & cell_state_fwd, const at::Tensor & input, const at::Tensor & layersOutputs, at::TensorList hx, at::TensorList params, bool has_biases, int64_t num_layers, double dropout, bool train, bool bidirectional, bool batch_first) {
  14011. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14012. auto maybe_layer = maybeCurrentDynamicLayer();
  14013. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14014. int64_t cur_level = maybe_layer->layerId();
  14015. if (!isBatchedAtLevel(grad_y, cur_level) && !isBatchedAtLevel(grad_hy, cur_level) && !isBatchedAtLevel(grad_cy, cur_level) && !isBatchedAtLevel(z_state, cur_level) && !isBatchedAtLevel(cell_state_fwd, cur_level) && !isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(layersOutputs, cur_level) && !isBatchedAtLevel(hx, cur_level) && !isBatchedAtLevel(params, cur_level)) {
  14016. return at::_ops::lstm_mps_backward::call(grad_y, grad_hy, grad_cy, z_state, cell_state_fwd, input, layersOutputs, hx, params, has_biases, num_layers, dropout, train, bidirectional, batch_first);
  14017. }
  14018. auto [z_state_value, z_state_bdim] = unwrapTensorAtLevel(z_state, cur_level);
  14019. auto [cell_state_fwd_value, cell_state_fwd_bdim] = unwrapTensorAtLevel(cell_state_fwd, cur_level);
  14020. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  14021. auto [layersOutputs_value, layersOutputs_bdim] = unwrapTensorAtLevel(layersOutputs, cur_level);
  14022. std::optional<Tensor> grad_y_value;
  14023. std::optional<int64_t> grad_y_bdim;
  14024. if (grad_y) {
  14025. std::tie(grad_y_value, grad_y_bdim) = unwrapTensorAtLevel(grad_y.value(), cur_level);
  14026. }
  14027. std::optional<Tensor> grad_hy_value;
  14028. std::optional<int64_t> grad_hy_bdim;
  14029. if (grad_hy) {
  14030. std::tie(grad_hy_value, grad_hy_bdim) = unwrapTensorAtLevel(grad_hy.value(), cur_level);
  14031. }
  14032. std::optional<Tensor> grad_cy_value;
  14033. std::optional<int64_t> grad_cy_bdim;
  14034. if (grad_cy) {
  14035. std::tie(grad_cy_value, grad_cy_bdim) = unwrapTensorAtLevel(grad_cy.value(), cur_level);
  14036. }
  14037. auto results = batch_rule(grad_y_value, grad_y_bdim, grad_hy_value, grad_hy_bdim, grad_cy_value, grad_cy_bdim, z_state_value, z_state_bdim, cell_state_fwd_value, cell_state_fwd_bdim, input_value, input_bdim, layersOutputs_value, layersOutputs_bdim, hx, params, has_biases, num_layers, dropout, train, bidirectional, batch_first);
  14038. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatchedVector(std::get<2>(results), std::get<3>(results), cur_level), makeBatchedVector(std::get<4>(results), std::get<5>(results), cur_level));
  14039. }
  14040. template <typename batch_rule_t, batch_rule_t batch_rule>
  14041. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> _thnn_fused_lstm_cell_generated_plumbing(const at::Tensor & input_gates, const at::Tensor & hidden_gates, const at::Tensor & cx, const ::std::optional<at::Tensor> & input_bias, const ::std::optional<at::Tensor> & hidden_bias) {
  14042. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14043. auto maybe_layer = maybeCurrentDynamicLayer();
  14044. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14045. int64_t cur_level = maybe_layer->layerId();
  14046. if (!isBatchedAtLevel(input_gates, cur_level) && !isBatchedAtLevel(hidden_gates, cur_level) && !isBatchedAtLevel(cx, cur_level) && !isBatchedAtLevel(input_bias, cur_level) && !isBatchedAtLevel(hidden_bias, cur_level)) {
  14047. return at::_ops::_thnn_fused_lstm_cell::call(input_gates, hidden_gates, cx, input_bias, hidden_bias);
  14048. }
  14049. auto [input_gates_value, input_gates_bdim] = unwrapTensorAtLevel(input_gates, cur_level);
  14050. auto [hidden_gates_value, hidden_gates_bdim] = unwrapTensorAtLevel(hidden_gates, cur_level);
  14051. auto [cx_value, cx_bdim] = unwrapTensorAtLevel(cx, cur_level);
  14052. std::optional<Tensor> input_bias_value;
  14053. std::optional<int64_t> input_bias_bdim;
  14054. if (input_bias) {
  14055. std::tie(input_bias_value, input_bias_bdim) = unwrapTensorAtLevel(input_bias.value(), cur_level);
  14056. }
  14057. std::optional<Tensor> hidden_bias_value;
  14058. std::optional<int64_t> hidden_bias_bdim;
  14059. if (hidden_bias) {
  14060. std::tie(hidden_bias_value, hidden_bias_bdim) = unwrapTensorAtLevel(hidden_bias.value(), cur_level);
  14061. }
  14062. auto results = batch_rule(input_gates_value, input_gates_bdim, hidden_gates_value, hidden_gates_bdim, cx_value, cx_bdim, input_bias_value, input_bias_bdim, hidden_bias_value, hidden_bias_bdim);
  14063. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  14064. }
  14065. template <typename batch_rule_t, batch_rule_t batch_rule>
  14066. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> _thnn_fused_lstm_cell_backward_impl_generated_plumbing(const ::std::optional<at::Tensor> & grad_hy, const ::std::optional<at::Tensor> & grad_cy, const at::Tensor & cx, const at::Tensor & cy, const at::Tensor & workspace, bool has_bias) {
  14067. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14068. auto maybe_layer = maybeCurrentDynamicLayer();
  14069. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14070. int64_t cur_level = maybe_layer->layerId();
  14071. if (!isBatchedAtLevel(grad_hy, cur_level) && !isBatchedAtLevel(grad_cy, cur_level) && !isBatchedAtLevel(cx, cur_level) && !isBatchedAtLevel(cy, cur_level) && !isBatchedAtLevel(workspace, cur_level)) {
  14072. return at::_ops::_thnn_fused_lstm_cell_backward_impl::call(grad_hy, grad_cy, cx, cy, workspace, has_bias);
  14073. }
  14074. auto [cx_value, cx_bdim] = unwrapTensorAtLevel(cx, cur_level);
  14075. auto [cy_value, cy_bdim] = unwrapTensorAtLevel(cy, cur_level);
  14076. auto [workspace_value, workspace_bdim] = unwrapTensorAtLevel(workspace, cur_level);
  14077. std::optional<Tensor> grad_hy_value;
  14078. std::optional<int64_t> grad_hy_bdim;
  14079. if (grad_hy) {
  14080. std::tie(grad_hy_value, grad_hy_bdim) = unwrapTensorAtLevel(grad_hy.value(), cur_level);
  14081. }
  14082. std::optional<Tensor> grad_cy_value;
  14083. std::optional<int64_t> grad_cy_bdim;
  14084. if (grad_cy) {
  14085. std::tie(grad_cy_value, grad_cy_bdim) = unwrapTensorAtLevel(grad_cy.value(), cur_level);
  14086. }
  14087. auto results = batch_rule(grad_hy_value, grad_hy_bdim, grad_cy_value, grad_cy_bdim, cx_value, cx_bdim, cy_value, cy_bdim, workspace_value, workspace_bdim, has_bias);
  14088. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  14089. }
  14090. template <typename batch_rule_t, batch_rule_t batch_rule>
  14091. ::std::tuple<at::Tensor,at::Tensor,at::Tensor,at::Tensor,at::Tensor> _thnn_fused_lstm_cell_backward_generated_plumbing(const ::std::optional<at::Tensor> & grad_hy, const ::std::optional<at::Tensor> & grad_cy, const at::Tensor & cx, const at::Tensor & cy, const at::Tensor & workspace, bool has_bias) {
  14092. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14093. auto maybe_layer = maybeCurrentDynamicLayer();
  14094. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14095. int64_t cur_level = maybe_layer->layerId();
  14096. if (!isBatchedAtLevel(grad_hy, cur_level) && !isBatchedAtLevel(grad_cy, cur_level) && !isBatchedAtLevel(cx, cur_level) && !isBatchedAtLevel(cy, cur_level) && !isBatchedAtLevel(workspace, cur_level)) {
  14097. return at::_ops::_thnn_fused_lstm_cell_backward::call(grad_hy, grad_cy, cx, cy, workspace, has_bias);
  14098. }
  14099. auto [cx_value, cx_bdim] = unwrapTensorAtLevel(cx, cur_level);
  14100. auto [cy_value, cy_bdim] = unwrapTensorAtLevel(cy, cur_level);
  14101. auto [workspace_value, workspace_bdim] = unwrapTensorAtLevel(workspace, cur_level);
  14102. std::optional<Tensor> grad_hy_value;
  14103. std::optional<int64_t> grad_hy_bdim;
  14104. if (grad_hy) {
  14105. std::tie(grad_hy_value, grad_hy_bdim) = unwrapTensorAtLevel(grad_hy.value(), cur_level);
  14106. }
  14107. std::optional<Tensor> grad_cy_value;
  14108. std::optional<int64_t> grad_cy_bdim;
  14109. if (grad_cy) {
  14110. std::tie(grad_cy_value, grad_cy_bdim) = unwrapTensorAtLevel(grad_cy.value(), cur_level);
  14111. }
  14112. auto results = batch_rule(grad_hy_value, grad_hy_bdim, grad_cy_value, grad_cy_bdim, cx_value, cx_bdim, cy_value, cy_bdim, workspace_value, workspace_bdim, has_bias);
  14113. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level), makeBatched(std::get<6>(results), std::get<7>(results), cur_level), makeBatched(std::get<8>(results), std::get<9>(results), cur_level));
  14114. }
  14115. template <typename batch_rule_t, batch_rule_t batch_rule>
  14116. ::std::tuple<at::Tensor,at::Tensor,at::Tensor,at::Tensor,at::Tensor> _thnn_differentiable_lstm_cell_backward_generated_plumbing(const ::std::optional<at::Tensor> & grad_hy, const ::std::optional<at::Tensor> & grad_cy, const at::Tensor & input_gates, const at::Tensor & hidden_gates, const ::std::optional<at::Tensor> & input_bias, const ::std::optional<at::Tensor> & hidden_bias, const at::Tensor & cx, const at::Tensor & cy) {
  14117. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14118. auto maybe_layer = maybeCurrentDynamicLayer();
  14119. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14120. int64_t cur_level = maybe_layer->layerId();
  14121. if (!isBatchedAtLevel(grad_hy, cur_level) && !isBatchedAtLevel(grad_cy, cur_level) && !isBatchedAtLevel(input_gates, cur_level) && !isBatchedAtLevel(hidden_gates, cur_level) && !isBatchedAtLevel(input_bias, cur_level) && !isBatchedAtLevel(hidden_bias, cur_level) && !isBatchedAtLevel(cx, cur_level) && !isBatchedAtLevel(cy, cur_level)) {
  14122. return at::_ops::_thnn_differentiable_lstm_cell_backward::call(grad_hy, grad_cy, input_gates, hidden_gates, input_bias, hidden_bias, cx, cy);
  14123. }
  14124. auto [input_gates_value, input_gates_bdim] = unwrapTensorAtLevel(input_gates, cur_level);
  14125. auto [hidden_gates_value, hidden_gates_bdim] = unwrapTensorAtLevel(hidden_gates, cur_level);
  14126. auto [cx_value, cx_bdim] = unwrapTensorAtLevel(cx, cur_level);
  14127. auto [cy_value, cy_bdim] = unwrapTensorAtLevel(cy, cur_level);
  14128. std::optional<Tensor> grad_hy_value;
  14129. std::optional<int64_t> grad_hy_bdim;
  14130. if (grad_hy) {
  14131. std::tie(grad_hy_value, grad_hy_bdim) = unwrapTensorAtLevel(grad_hy.value(), cur_level);
  14132. }
  14133. std::optional<Tensor> grad_cy_value;
  14134. std::optional<int64_t> grad_cy_bdim;
  14135. if (grad_cy) {
  14136. std::tie(grad_cy_value, grad_cy_bdim) = unwrapTensorAtLevel(grad_cy.value(), cur_level);
  14137. }
  14138. std::optional<Tensor> input_bias_value;
  14139. std::optional<int64_t> input_bias_bdim;
  14140. if (input_bias) {
  14141. std::tie(input_bias_value, input_bias_bdim) = unwrapTensorAtLevel(input_bias.value(), cur_level);
  14142. }
  14143. std::optional<Tensor> hidden_bias_value;
  14144. std::optional<int64_t> hidden_bias_bdim;
  14145. if (hidden_bias) {
  14146. std::tie(hidden_bias_value, hidden_bias_bdim) = unwrapTensorAtLevel(hidden_bias.value(), cur_level);
  14147. }
  14148. auto results = batch_rule(grad_hy_value, grad_hy_bdim, grad_cy_value, grad_cy_bdim, input_gates_value, input_gates_bdim, hidden_gates_value, hidden_gates_bdim, input_bias_value, input_bias_bdim, hidden_bias_value, hidden_bias_bdim, cx_value, cx_bdim, cy_value, cy_bdim);
  14149. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level), makeBatched(std::get<6>(results), std::get<7>(results), cur_level), makeBatched(std::get<8>(results), std::get<9>(results), cur_level));
  14150. }
  14151. template <typename batch_rule_t, batch_rule_t batch_rule>
  14152. ::std::tuple<at::Tensor,at::Tensor> _thnn_fused_gru_cell_generated_plumbing(const at::Tensor & input_gates, const at::Tensor & hidden_gates, const at::Tensor & hx, const ::std::optional<at::Tensor> & input_bias, const ::std::optional<at::Tensor> & hidden_bias) {
  14153. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14154. auto maybe_layer = maybeCurrentDynamicLayer();
  14155. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14156. int64_t cur_level = maybe_layer->layerId();
  14157. if (!isBatchedAtLevel(input_gates, cur_level) && !isBatchedAtLevel(hidden_gates, cur_level) && !isBatchedAtLevel(hx, cur_level) && !isBatchedAtLevel(input_bias, cur_level) && !isBatchedAtLevel(hidden_bias, cur_level)) {
  14158. return at::_ops::_thnn_fused_gru_cell::call(input_gates, hidden_gates, hx, input_bias, hidden_bias);
  14159. }
  14160. auto [input_gates_value, input_gates_bdim] = unwrapTensorAtLevel(input_gates, cur_level);
  14161. auto [hidden_gates_value, hidden_gates_bdim] = unwrapTensorAtLevel(hidden_gates, cur_level);
  14162. auto [hx_value, hx_bdim] = unwrapTensorAtLevel(hx, cur_level);
  14163. std::optional<Tensor> input_bias_value;
  14164. std::optional<int64_t> input_bias_bdim;
  14165. if (input_bias) {
  14166. std::tie(input_bias_value, input_bias_bdim) = unwrapTensorAtLevel(input_bias.value(), cur_level);
  14167. }
  14168. std::optional<Tensor> hidden_bias_value;
  14169. std::optional<int64_t> hidden_bias_bdim;
  14170. if (hidden_bias) {
  14171. std::tie(hidden_bias_value, hidden_bias_bdim) = unwrapTensorAtLevel(hidden_bias.value(), cur_level);
  14172. }
  14173. auto results = batch_rule(input_gates_value, input_gates_bdim, hidden_gates_value, hidden_gates_bdim, hx_value, hx_bdim, input_bias_value, input_bias_bdim, hidden_bias_value, hidden_bias_bdim);
  14174. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  14175. }
  14176. template <typename batch_rule_t, batch_rule_t batch_rule>
  14177. ::std::tuple<at::Tensor,at::Tensor,at::Tensor,at::Tensor,at::Tensor> _thnn_fused_gru_cell_backward_generated_plumbing(const at::Tensor & grad_hy, const at::Tensor & workspace, bool has_bias) {
  14178. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14179. auto maybe_layer = maybeCurrentDynamicLayer();
  14180. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14181. int64_t cur_level = maybe_layer->layerId();
  14182. if (!isBatchedAtLevel(grad_hy, cur_level) && !isBatchedAtLevel(workspace, cur_level)) {
  14183. return at::_ops::_thnn_fused_gru_cell_backward::call(grad_hy, workspace, has_bias);
  14184. }
  14185. auto [grad_hy_value, grad_hy_bdim] = unwrapTensorAtLevel(grad_hy, cur_level);
  14186. auto [workspace_value, workspace_bdim] = unwrapTensorAtLevel(workspace, cur_level);
  14187. auto results = batch_rule(grad_hy_value, grad_hy_bdim, workspace_value, workspace_bdim, has_bias);
  14188. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level), makeBatched(std::get<6>(results), std::get<7>(results), cur_level), makeBatched(std::get<8>(results), std::get<9>(results), cur_level));
  14189. }
  14190. template <typename batch_rule_t, batch_rule_t batch_rule>
  14191. ::std::tuple<at::Tensor,at::Tensor,at::Tensor,at::Tensor,at::Tensor> _thnn_differentiable_gru_cell_backward_generated_plumbing(const at::Tensor & grad_hy, const at::Tensor & input_gates, const at::Tensor & hidden_gates, const at::Tensor & hx, const ::std::optional<at::Tensor> & input_bias, const ::std::optional<at::Tensor> & hidden_bias) {
  14192. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14193. auto maybe_layer = maybeCurrentDynamicLayer();
  14194. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14195. int64_t cur_level = maybe_layer->layerId();
  14196. if (!isBatchedAtLevel(grad_hy, cur_level) && !isBatchedAtLevel(input_gates, cur_level) && !isBatchedAtLevel(hidden_gates, cur_level) && !isBatchedAtLevel(hx, cur_level) && !isBatchedAtLevel(input_bias, cur_level) && !isBatchedAtLevel(hidden_bias, cur_level)) {
  14197. return at::_ops::_thnn_differentiable_gru_cell_backward::call(grad_hy, input_gates, hidden_gates, hx, input_bias, hidden_bias);
  14198. }
  14199. auto [grad_hy_value, grad_hy_bdim] = unwrapTensorAtLevel(grad_hy, cur_level);
  14200. auto [input_gates_value, input_gates_bdim] = unwrapTensorAtLevel(input_gates, cur_level);
  14201. auto [hidden_gates_value, hidden_gates_bdim] = unwrapTensorAtLevel(hidden_gates, cur_level);
  14202. auto [hx_value, hx_bdim] = unwrapTensorAtLevel(hx, cur_level);
  14203. std::optional<Tensor> input_bias_value;
  14204. std::optional<int64_t> input_bias_bdim;
  14205. if (input_bias) {
  14206. std::tie(input_bias_value, input_bias_bdim) = unwrapTensorAtLevel(input_bias.value(), cur_level);
  14207. }
  14208. std::optional<Tensor> hidden_bias_value;
  14209. std::optional<int64_t> hidden_bias_bdim;
  14210. if (hidden_bias) {
  14211. std::tie(hidden_bias_value, hidden_bias_bdim) = unwrapTensorAtLevel(hidden_bias.value(), cur_level);
  14212. }
  14213. auto results = batch_rule(grad_hy_value, grad_hy_bdim, input_gates_value, input_gates_bdim, hidden_gates_value, hidden_gates_bdim, hx_value, hx_bdim, input_bias_value, input_bias_bdim, hidden_bias_value, hidden_bias_bdim);
  14214. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level), makeBatched(std::get<6>(results), std::get<7>(results), cur_level), makeBatched(std::get<8>(results), std::get<9>(results), cur_level));
  14215. }
  14216. template <typename batch_rule_t, batch_rule_t batch_rule>
  14217. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> lstm_input_generated_plumbing(const at::Tensor & input, at::TensorList hx, at::TensorList params, bool has_biases, int64_t num_layers, double dropout, bool train, bool bidirectional, bool batch_first) {
  14218. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14219. auto maybe_layer = maybeCurrentDynamicLayer();
  14220. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14221. int64_t cur_level = maybe_layer->layerId();
  14222. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(hx, cur_level) && !isBatchedAtLevel(params, cur_level)) {
  14223. return at::_ops::lstm_input::call(input, hx, params, has_biases, num_layers, dropout, train, bidirectional, batch_first);
  14224. }
  14225. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  14226. auto results = batch_rule(input_value, input_bdim, hx, params, has_biases, num_layers, dropout, train, bidirectional, batch_first);
  14227. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  14228. }
  14229. template <typename batch_rule_t, batch_rule_t batch_rule>
  14230. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> lstm_data_generated_plumbing(const at::Tensor & data, const at::Tensor & batch_sizes, at::TensorList hx, at::TensorList params, bool has_biases, int64_t num_layers, double dropout, bool train, bool bidirectional) {
  14231. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14232. auto maybe_layer = maybeCurrentDynamicLayer();
  14233. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14234. int64_t cur_level = maybe_layer->layerId();
  14235. if (!isBatchedAtLevel(data, cur_level) && !isBatchedAtLevel(batch_sizes, cur_level) && !isBatchedAtLevel(hx, cur_level) && !isBatchedAtLevel(params, cur_level)) {
  14236. return at::_ops::lstm_data::call(data, batch_sizes, hx, params, has_biases, num_layers, dropout, train, bidirectional);
  14237. }
  14238. auto [data_value, data_bdim] = unwrapTensorAtLevel(data, cur_level);
  14239. auto [batch_sizes_value, batch_sizes_bdim] = unwrapTensorAtLevel(batch_sizes, cur_level);
  14240. auto results = batch_rule(data_value, data_bdim, batch_sizes_value, batch_sizes_bdim, hx, params, has_biases, num_layers, dropout, train, bidirectional);
  14241. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  14242. }
  14243. template <typename batch_rule_t, batch_rule_t batch_rule>
  14244. ::std::tuple<at::Tensor,at::Tensor> gru_input_generated_plumbing(const at::Tensor & input, const at::Tensor & hx, at::TensorList params, bool has_biases, int64_t num_layers, double dropout, bool train, bool bidirectional, bool batch_first) {
  14245. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14246. auto maybe_layer = maybeCurrentDynamicLayer();
  14247. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14248. int64_t cur_level = maybe_layer->layerId();
  14249. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(hx, cur_level) && !isBatchedAtLevel(params, cur_level)) {
  14250. return at::_ops::gru_input::call(input, hx, params, has_biases, num_layers, dropout, train, bidirectional, batch_first);
  14251. }
  14252. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  14253. auto [hx_value, hx_bdim] = unwrapTensorAtLevel(hx, cur_level);
  14254. auto results = batch_rule(input_value, input_bdim, hx_value, hx_bdim, params, has_biases, num_layers, dropout, train, bidirectional, batch_first);
  14255. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  14256. }
  14257. template <typename batch_rule_t, batch_rule_t batch_rule>
  14258. ::std::tuple<at::Tensor,at::Tensor> gru_data_generated_plumbing(const at::Tensor & data, const at::Tensor & batch_sizes, const at::Tensor & hx, at::TensorList params, bool has_biases, int64_t num_layers, double dropout, bool train, bool bidirectional) {
  14259. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14260. auto maybe_layer = maybeCurrentDynamicLayer();
  14261. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14262. int64_t cur_level = maybe_layer->layerId();
  14263. if (!isBatchedAtLevel(data, cur_level) && !isBatchedAtLevel(batch_sizes, cur_level) && !isBatchedAtLevel(hx, cur_level) && !isBatchedAtLevel(params, cur_level)) {
  14264. return at::_ops::gru_data::call(data, batch_sizes, hx, params, has_biases, num_layers, dropout, train, bidirectional);
  14265. }
  14266. auto [data_value, data_bdim] = unwrapTensorAtLevel(data, cur_level);
  14267. auto [batch_sizes_value, batch_sizes_bdim] = unwrapTensorAtLevel(batch_sizes, cur_level);
  14268. auto [hx_value, hx_bdim] = unwrapTensorAtLevel(hx, cur_level);
  14269. auto results = batch_rule(data_value, data_bdim, batch_sizes_value, batch_sizes_bdim, hx_value, hx_bdim, params, has_biases, num_layers, dropout, train, bidirectional);
  14270. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  14271. }
  14272. template <typename batch_rule_t, batch_rule_t batch_rule>
  14273. ::std::tuple<at::Tensor,at::Tensor> rnn_tanh_input_generated_plumbing(const at::Tensor & input, const at::Tensor & hx, at::TensorList params, bool has_biases, int64_t num_layers, double dropout, bool train, bool bidirectional, bool batch_first) {
  14274. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14275. auto maybe_layer = maybeCurrentDynamicLayer();
  14276. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14277. int64_t cur_level = maybe_layer->layerId();
  14278. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(hx, cur_level) && !isBatchedAtLevel(params, cur_level)) {
  14279. return at::_ops::rnn_tanh_input::call(input, hx, params, has_biases, num_layers, dropout, train, bidirectional, batch_first);
  14280. }
  14281. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  14282. auto [hx_value, hx_bdim] = unwrapTensorAtLevel(hx, cur_level);
  14283. auto results = batch_rule(input_value, input_bdim, hx_value, hx_bdim, params, has_biases, num_layers, dropout, train, bidirectional, batch_first);
  14284. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  14285. }
  14286. template <typename batch_rule_t, batch_rule_t batch_rule>
  14287. ::std::tuple<at::Tensor,at::Tensor> rnn_tanh_data_generated_plumbing(const at::Tensor & data, const at::Tensor & batch_sizes, const at::Tensor & hx, at::TensorList params, bool has_biases, int64_t num_layers, double dropout, bool train, bool bidirectional) {
  14288. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14289. auto maybe_layer = maybeCurrentDynamicLayer();
  14290. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14291. int64_t cur_level = maybe_layer->layerId();
  14292. if (!isBatchedAtLevel(data, cur_level) && !isBatchedAtLevel(batch_sizes, cur_level) && !isBatchedAtLevel(hx, cur_level) && !isBatchedAtLevel(params, cur_level)) {
  14293. return at::_ops::rnn_tanh_data::call(data, batch_sizes, hx, params, has_biases, num_layers, dropout, train, bidirectional);
  14294. }
  14295. auto [data_value, data_bdim] = unwrapTensorAtLevel(data, cur_level);
  14296. auto [batch_sizes_value, batch_sizes_bdim] = unwrapTensorAtLevel(batch_sizes, cur_level);
  14297. auto [hx_value, hx_bdim] = unwrapTensorAtLevel(hx, cur_level);
  14298. auto results = batch_rule(data_value, data_bdim, batch_sizes_value, batch_sizes_bdim, hx_value, hx_bdim, params, has_biases, num_layers, dropout, train, bidirectional);
  14299. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  14300. }
  14301. template <typename batch_rule_t, batch_rule_t batch_rule>
  14302. ::std::tuple<at::Tensor,at::Tensor> rnn_relu_input_generated_plumbing(const at::Tensor & input, const at::Tensor & hx, at::TensorList params, bool has_biases, int64_t num_layers, double dropout, bool train, bool bidirectional, bool batch_first) {
  14303. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14304. auto maybe_layer = maybeCurrentDynamicLayer();
  14305. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14306. int64_t cur_level = maybe_layer->layerId();
  14307. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(hx, cur_level) && !isBatchedAtLevel(params, cur_level)) {
  14308. return at::_ops::rnn_relu_input::call(input, hx, params, has_biases, num_layers, dropout, train, bidirectional, batch_first);
  14309. }
  14310. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  14311. auto [hx_value, hx_bdim] = unwrapTensorAtLevel(hx, cur_level);
  14312. auto results = batch_rule(input_value, input_bdim, hx_value, hx_bdim, params, has_biases, num_layers, dropout, train, bidirectional, batch_first);
  14313. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  14314. }
  14315. template <typename batch_rule_t, batch_rule_t batch_rule>
  14316. ::std::tuple<at::Tensor,at::Tensor> rnn_relu_data_generated_plumbing(const at::Tensor & data, const at::Tensor & batch_sizes, const at::Tensor & hx, at::TensorList params, bool has_biases, int64_t num_layers, double dropout, bool train, bool bidirectional) {
  14317. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14318. auto maybe_layer = maybeCurrentDynamicLayer();
  14319. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14320. int64_t cur_level = maybe_layer->layerId();
  14321. if (!isBatchedAtLevel(data, cur_level) && !isBatchedAtLevel(batch_sizes, cur_level) && !isBatchedAtLevel(hx, cur_level) && !isBatchedAtLevel(params, cur_level)) {
  14322. return at::_ops::rnn_relu_data::call(data, batch_sizes, hx, params, has_biases, num_layers, dropout, train, bidirectional);
  14323. }
  14324. auto [data_value, data_bdim] = unwrapTensorAtLevel(data, cur_level);
  14325. auto [batch_sizes_value, batch_sizes_bdim] = unwrapTensorAtLevel(batch_sizes, cur_level);
  14326. auto [hx_value, hx_bdim] = unwrapTensorAtLevel(hx, cur_level);
  14327. auto results = batch_rule(data_value, data_bdim, batch_sizes_value, batch_sizes_bdim, hx_value, hx_bdim, params, has_biases, num_layers, dropout, train, bidirectional);
  14328. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  14329. }
  14330. template <typename batch_rule_t, batch_rule_t batch_rule>
  14331. ::std::tuple<at::Tensor,at::Tensor> lstm_cell_generated_plumbing(const at::Tensor & input, at::TensorList hx, const at::Tensor & w_ih, const at::Tensor & w_hh, const ::std::optional<at::Tensor> & b_ih, const ::std::optional<at::Tensor> & b_hh) {
  14332. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14333. auto maybe_layer = maybeCurrentDynamicLayer();
  14334. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14335. int64_t cur_level = maybe_layer->layerId();
  14336. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(hx, cur_level) && !isBatchedAtLevel(w_ih, cur_level) && !isBatchedAtLevel(w_hh, cur_level) && !isBatchedAtLevel(b_ih, cur_level) && !isBatchedAtLevel(b_hh, cur_level)) {
  14337. return at::_ops::lstm_cell::call(input, hx, w_ih, w_hh, b_ih, b_hh);
  14338. }
  14339. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  14340. auto [w_ih_value, w_ih_bdim] = unwrapTensorAtLevel(w_ih, cur_level);
  14341. auto [w_hh_value, w_hh_bdim] = unwrapTensorAtLevel(w_hh, cur_level);
  14342. std::optional<Tensor> b_ih_value;
  14343. std::optional<int64_t> b_ih_bdim;
  14344. if (b_ih) {
  14345. std::tie(b_ih_value, b_ih_bdim) = unwrapTensorAtLevel(b_ih.value(), cur_level);
  14346. }
  14347. std::optional<Tensor> b_hh_value;
  14348. std::optional<int64_t> b_hh_bdim;
  14349. if (b_hh) {
  14350. std::tie(b_hh_value, b_hh_bdim) = unwrapTensorAtLevel(b_hh.value(), cur_level);
  14351. }
  14352. auto results = batch_rule(input_value, input_bdim, hx, w_ih_value, w_ih_bdim, w_hh_value, w_hh_bdim, b_ih_value, b_ih_bdim, b_hh_value, b_hh_bdim);
  14353. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  14354. }
  14355. template <typename batch_rule_t, batch_rule_t batch_rule>
  14356. at::Tensor gru_cell_generated_plumbing(const at::Tensor & input, const at::Tensor & hx, const at::Tensor & w_ih, const at::Tensor & w_hh, const ::std::optional<at::Tensor> & b_ih, const ::std::optional<at::Tensor> & b_hh) {
  14357. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14358. auto maybe_layer = maybeCurrentDynamicLayer();
  14359. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14360. int64_t cur_level = maybe_layer->layerId();
  14361. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(hx, cur_level) && !isBatchedAtLevel(w_ih, cur_level) && !isBatchedAtLevel(w_hh, cur_level) && !isBatchedAtLevel(b_ih, cur_level) && !isBatchedAtLevel(b_hh, cur_level)) {
  14362. return at::_ops::gru_cell::call(input, hx, w_ih, w_hh, b_ih, b_hh);
  14363. }
  14364. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  14365. auto [hx_value, hx_bdim] = unwrapTensorAtLevel(hx, cur_level);
  14366. auto [w_ih_value, w_ih_bdim] = unwrapTensorAtLevel(w_ih, cur_level);
  14367. auto [w_hh_value, w_hh_bdim] = unwrapTensorAtLevel(w_hh, cur_level);
  14368. std::optional<Tensor> b_ih_value;
  14369. std::optional<int64_t> b_ih_bdim;
  14370. if (b_ih) {
  14371. std::tie(b_ih_value, b_ih_bdim) = unwrapTensorAtLevel(b_ih.value(), cur_level);
  14372. }
  14373. std::optional<Tensor> b_hh_value;
  14374. std::optional<int64_t> b_hh_bdim;
  14375. if (b_hh) {
  14376. std::tie(b_hh_value, b_hh_bdim) = unwrapTensorAtLevel(b_hh.value(), cur_level);
  14377. }
  14378. auto results = batch_rule(input_value, input_bdim, hx_value, hx_bdim, w_ih_value, w_ih_bdim, w_hh_value, w_hh_bdim, b_ih_value, b_ih_bdim, b_hh_value, b_hh_bdim);
  14379. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  14380. }
  14381. template <typename batch_rule_t, batch_rule_t batch_rule>
  14382. at::Tensor rnn_tanh_cell_generated_plumbing(const at::Tensor & input, const at::Tensor & hx, const at::Tensor & w_ih, const at::Tensor & w_hh, const ::std::optional<at::Tensor> & b_ih, const ::std::optional<at::Tensor> & b_hh) {
  14383. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14384. auto maybe_layer = maybeCurrentDynamicLayer();
  14385. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14386. int64_t cur_level = maybe_layer->layerId();
  14387. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(hx, cur_level) && !isBatchedAtLevel(w_ih, cur_level) && !isBatchedAtLevel(w_hh, cur_level) && !isBatchedAtLevel(b_ih, cur_level) && !isBatchedAtLevel(b_hh, cur_level)) {
  14388. return at::_ops::rnn_tanh_cell::call(input, hx, w_ih, w_hh, b_ih, b_hh);
  14389. }
  14390. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  14391. auto [hx_value, hx_bdim] = unwrapTensorAtLevel(hx, cur_level);
  14392. auto [w_ih_value, w_ih_bdim] = unwrapTensorAtLevel(w_ih, cur_level);
  14393. auto [w_hh_value, w_hh_bdim] = unwrapTensorAtLevel(w_hh, cur_level);
  14394. std::optional<Tensor> b_ih_value;
  14395. std::optional<int64_t> b_ih_bdim;
  14396. if (b_ih) {
  14397. std::tie(b_ih_value, b_ih_bdim) = unwrapTensorAtLevel(b_ih.value(), cur_level);
  14398. }
  14399. std::optional<Tensor> b_hh_value;
  14400. std::optional<int64_t> b_hh_bdim;
  14401. if (b_hh) {
  14402. std::tie(b_hh_value, b_hh_bdim) = unwrapTensorAtLevel(b_hh.value(), cur_level);
  14403. }
  14404. auto results = batch_rule(input_value, input_bdim, hx_value, hx_bdim, w_ih_value, w_ih_bdim, w_hh_value, w_hh_bdim, b_ih_value, b_ih_bdim, b_hh_value, b_hh_bdim);
  14405. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  14406. }
  14407. template <typename batch_rule_t, batch_rule_t batch_rule>
  14408. at::Tensor rnn_relu_cell_generated_plumbing(const at::Tensor & input, const at::Tensor & hx, const at::Tensor & w_ih, const at::Tensor & w_hh, const ::std::optional<at::Tensor> & b_ih, const ::std::optional<at::Tensor> & b_hh) {
  14409. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14410. auto maybe_layer = maybeCurrentDynamicLayer();
  14411. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14412. int64_t cur_level = maybe_layer->layerId();
  14413. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(hx, cur_level) && !isBatchedAtLevel(w_ih, cur_level) && !isBatchedAtLevel(w_hh, cur_level) && !isBatchedAtLevel(b_ih, cur_level) && !isBatchedAtLevel(b_hh, cur_level)) {
  14414. return at::_ops::rnn_relu_cell::call(input, hx, w_ih, w_hh, b_ih, b_hh);
  14415. }
  14416. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  14417. auto [hx_value, hx_bdim] = unwrapTensorAtLevel(hx, cur_level);
  14418. auto [w_ih_value, w_ih_bdim] = unwrapTensorAtLevel(w_ih, cur_level);
  14419. auto [w_hh_value, w_hh_bdim] = unwrapTensorAtLevel(w_hh, cur_level);
  14420. std::optional<Tensor> b_ih_value;
  14421. std::optional<int64_t> b_ih_bdim;
  14422. if (b_ih) {
  14423. std::tie(b_ih_value, b_ih_bdim) = unwrapTensorAtLevel(b_ih.value(), cur_level);
  14424. }
  14425. std::optional<Tensor> b_hh_value;
  14426. std::optional<int64_t> b_hh_bdim;
  14427. if (b_hh) {
  14428. std::tie(b_hh_value, b_hh_bdim) = unwrapTensorAtLevel(b_hh.value(), cur_level);
  14429. }
  14430. auto results = batch_rule(input_value, input_bdim, hx_value, hx_bdim, w_ih_value, w_ih_bdim, w_hh_value, w_hh_bdim, b_ih_value, b_ih_bdim, b_hh_value, b_hh_bdim);
  14431. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  14432. }
  14433. template <typename batch_rule_t, batch_rule_t batch_rule>
  14434. ::std::tuple<at::Tensor,at::Tensor> quantized_lstm_cell_generated_plumbing(const at::Tensor & input, at::TensorList hx, const at::Tensor & w_ih, const at::Tensor & w_hh, const at::Tensor & b_ih, const at::Tensor & b_hh, const at::Tensor & packed_ih, const at::Tensor & packed_hh, const at::Tensor & col_offsets_ih, const at::Tensor & col_offsets_hh, const at::Scalar & scale_ih, const at::Scalar & scale_hh, const at::Scalar & zero_point_ih, const at::Scalar & zero_point_hh) {
  14435. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14436. auto maybe_layer = maybeCurrentDynamicLayer();
  14437. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14438. int64_t cur_level = maybe_layer->layerId();
  14439. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(hx, cur_level) && !isBatchedAtLevel(w_ih, cur_level) && !isBatchedAtLevel(w_hh, cur_level) && !isBatchedAtLevel(b_ih, cur_level) && !isBatchedAtLevel(b_hh, cur_level) && !isBatchedAtLevel(packed_ih, cur_level) && !isBatchedAtLevel(packed_hh, cur_level) && !isBatchedAtLevel(col_offsets_ih, cur_level) && !isBatchedAtLevel(col_offsets_hh, cur_level)) {
  14440. return at::_ops::quantized_lstm_cell::call(input, hx, w_ih, w_hh, b_ih, b_hh, packed_ih, packed_hh, col_offsets_ih, col_offsets_hh, scale_ih, scale_hh, zero_point_ih, zero_point_hh);
  14441. }
  14442. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  14443. auto [w_ih_value, w_ih_bdim] = unwrapTensorAtLevel(w_ih, cur_level);
  14444. auto [w_hh_value, w_hh_bdim] = unwrapTensorAtLevel(w_hh, cur_level);
  14445. auto [b_ih_value, b_ih_bdim] = unwrapTensorAtLevel(b_ih, cur_level);
  14446. auto [b_hh_value, b_hh_bdim] = unwrapTensorAtLevel(b_hh, cur_level);
  14447. auto [packed_ih_value, packed_ih_bdim] = unwrapTensorAtLevel(packed_ih, cur_level);
  14448. auto [packed_hh_value, packed_hh_bdim] = unwrapTensorAtLevel(packed_hh, cur_level);
  14449. auto [col_offsets_ih_value, col_offsets_ih_bdim] = unwrapTensorAtLevel(col_offsets_ih, cur_level);
  14450. auto [col_offsets_hh_value, col_offsets_hh_bdim] = unwrapTensorAtLevel(col_offsets_hh, cur_level);
  14451. auto results = batch_rule(input_value, input_bdim, hx, w_ih_value, w_ih_bdim, w_hh_value, w_hh_bdim, b_ih_value, b_ih_bdim, b_hh_value, b_hh_bdim, packed_ih_value, packed_ih_bdim, packed_hh_value, packed_hh_bdim, col_offsets_ih_value, col_offsets_ih_bdim, col_offsets_hh_value, col_offsets_hh_bdim, scale_ih, scale_hh, zero_point_ih, zero_point_hh);
  14452. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  14453. }
  14454. template <typename batch_rule_t, batch_rule_t batch_rule>
  14455. at::Tensor quantized_gru_cell_generated_plumbing(const at::Tensor & input, const at::Tensor & hx, const at::Tensor & w_ih, const at::Tensor & w_hh, const at::Tensor & b_ih, const at::Tensor & b_hh, const at::Tensor & packed_ih, const at::Tensor & packed_hh, const at::Tensor & col_offsets_ih, const at::Tensor & col_offsets_hh, const at::Scalar & scale_ih, const at::Scalar & scale_hh, const at::Scalar & zero_point_ih, const at::Scalar & zero_point_hh) {
  14456. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14457. auto maybe_layer = maybeCurrentDynamicLayer();
  14458. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14459. int64_t cur_level = maybe_layer->layerId();
  14460. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(hx, cur_level) && !isBatchedAtLevel(w_ih, cur_level) && !isBatchedAtLevel(w_hh, cur_level) && !isBatchedAtLevel(b_ih, cur_level) && !isBatchedAtLevel(b_hh, cur_level) && !isBatchedAtLevel(packed_ih, cur_level) && !isBatchedAtLevel(packed_hh, cur_level) && !isBatchedAtLevel(col_offsets_ih, cur_level) && !isBatchedAtLevel(col_offsets_hh, cur_level)) {
  14461. return at::_ops::quantized_gru_cell::call(input, hx, w_ih, w_hh, b_ih, b_hh, packed_ih, packed_hh, col_offsets_ih, col_offsets_hh, scale_ih, scale_hh, zero_point_ih, zero_point_hh);
  14462. }
  14463. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  14464. auto [hx_value, hx_bdim] = unwrapTensorAtLevel(hx, cur_level);
  14465. auto [w_ih_value, w_ih_bdim] = unwrapTensorAtLevel(w_ih, cur_level);
  14466. auto [w_hh_value, w_hh_bdim] = unwrapTensorAtLevel(w_hh, cur_level);
  14467. auto [b_ih_value, b_ih_bdim] = unwrapTensorAtLevel(b_ih, cur_level);
  14468. auto [b_hh_value, b_hh_bdim] = unwrapTensorAtLevel(b_hh, cur_level);
  14469. auto [packed_ih_value, packed_ih_bdim] = unwrapTensorAtLevel(packed_ih, cur_level);
  14470. auto [packed_hh_value, packed_hh_bdim] = unwrapTensorAtLevel(packed_hh, cur_level);
  14471. auto [col_offsets_ih_value, col_offsets_ih_bdim] = unwrapTensorAtLevel(col_offsets_ih, cur_level);
  14472. auto [col_offsets_hh_value, col_offsets_hh_bdim] = unwrapTensorAtLevel(col_offsets_hh, cur_level);
  14473. auto results = batch_rule(input_value, input_bdim, hx_value, hx_bdim, w_ih_value, w_ih_bdim, w_hh_value, w_hh_bdim, b_ih_value, b_ih_bdim, b_hh_value, b_hh_bdim, packed_ih_value, packed_ih_bdim, packed_hh_value, packed_hh_bdim, col_offsets_ih_value, col_offsets_ih_bdim, col_offsets_hh_value, col_offsets_hh_bdim, scale_ih, scale_hh, zero_point_ih, zero_point_hh);
  14474. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  14475. }
  14476. template <typename batch_rule_t, batch_rule_t batch_rule>
  14477. at::Tensor quantized_rnn_relu_cell_generated_plumbing(const at::Tensor & input, const at::Tensor & hx, const at::Tensor & w_ih, const at::Tensor & w_hh, const at::Tensor & b_ih, const at::Tensor & b_hh, const at::Tensor & packed_ih, const at::Tensor & packed_hh, const at::Tensor & col_offsets_ih, const at::Tensor & col_offsets_hh, const at::Scalar & scale_ih, const at::Scalar & scale_hh, const at::Scalar & zero_point_ih, const at::Scalar & zero_point_hh) {
  14478. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14479. auto maybe_layer = maybeCurrentDynamicLayer();
  14480. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14481. int64_t cur_level = maybe_layer->layerId();
  14482. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(hx, cur_level) && !isBatchedAtLevel(w_ih, cur_level) && !isBatchedAtLevel(w_hh, cur_level) && !isBatchedAtLevel(b_ih, cur_level) && !isBatchedAtLevel(b_hh, cur_level) && !isBatchedAtLevel(packed_ih, cur_level) && !isBatchedAtLevel(packed_hh, cur_level) && !isBatchedAtLevel(col_offsets_ih, cur_level) && !isBatchedAtLevel(col_offsets_hh, cur_level)) {
  14483. return at::_ops::quantized_rnn_relu_cell::call(input, hx, w_ih, w_hh, b_ih, b_hh, packed_ih, packed_hh, col_offsets_ih, col_offsets_hh, scale_ih, scale_hh, zero_point_ih, zero_point_hh);
  14484. }
  14485. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  14486. auto [hx_value, hx_bdim] = unwrapTensorAtLevel(hx, cur_level);
  14487. auto [w_ih_value, w_ih_bdim] = unwrapTensorAtLevel(w_ih, cur_level);
  14488. auto [w_hh_value, w_hh_bdim] = unwrapTensorAtLevel(w_hh, cur_level);
  14489. auto [b_ih_value, b_ih_bdim] = unwrapTensorAtLevel(b_ih, cur_level);
  14490. auto [b_hh_value, b_hh_bdim] = unwrapTensorAtLevel(b_hh, cur_level);
  14491. auto [packed_ih_value, packed_ih_bdim] = unwrapTensorAtLevel(packed_ih, cur_level);
  14492. auto [packed_hh_value, packed_hh_bdim] = unwrapTensorAtLevel(packed_hh, cur_level);
  14493. auto [col_offsets_ih_value, col_offsets_ih_bdim] = unwrapTensorAtLevel(col_offsets_ih, cur_level);
  14494. auto [col_offsets_hh_value, col_offsets_hh_bdim] = unwrapTensorAtLevel(col_offsets_hh, cur_level);
  14495. auto results = batch_rule(input_value, input_bdim, hx_value, hx_bdim, w_ih_value, w_ih_bdim, w_hh_value, w_hh_bdim, b_ih_value, b_ih_bdim, b_hh_value, b_hh_bdim, packed_ih_value, packed_ih_bdim, packed_hh_value, packed_hh_bdim, col_offsets_ih_value, col_offsets_ih_bdim, col_offsets_hh_value, col_offsets_hh_bdim, scale_ih, scale_hh, zero_point_ih, zero_point_hh);
  14496. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  14497. }
  14498. template <typename batch_rule_t, batch_rule_t batch_rule>
  14499. at::Tensor quantized_rnn_tanh_cell_generated_plumbing(const at::Tensor & input, const at::Tensor & hx, const at::Tensor & w_ih, const at::Tensor & w_hh, const at::Tensor & b_ih, const at::Tensor & b_hh, const at::Tensor & packed_ih, const at::Tensor & packed_hh, const at::Tensor & col_offsets_ih, const at::Tensor & col_offsets_hh, const at::Scalar & scale_ih, const at::Scalar & scale_hh, const at::Scalar & zero_point_ih, const at::Scalar & zero_point_hh) {
  14500. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14501. auto maybe_layer = maybeCurrentDynamicLayer();
  14502. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14503. int64_t cur_level = maybe_layer->layerId();
  14504. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(hx, cur_level) && !isBatchedAtLevel(w_ih, cur_level) && !isBatchedAtLevel(w_hh, cur_level) && !isBatchedAtLevel(b_ih, cur_level) && !isBatchedAtLevel(b_hh, cur_level) && !isBatchedAtLevel(packed_ih, cur_level) && !isBatchedAtLevel(packed_hh, cur_level) && !isBatchedAtLevel(col_offsets_ih, cur_level) && !isBatchedAtLevel(col_offsets_hh, cur_level)) {
  14505. return at::_ops::quantized_rnn_tanh_cell::call(input, hx, w_ih, w_hh, b_ih, b_hh, packed_ih, packed_hh, col_offsets_ih, col_offsets_hh, scale_ih, scale_hh, zero_point_ih, zero_point_hh);
  14506. }
  14507. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  14508. auto [hx_value, hx_bdim] = unwrapTensorAtLevel(hx, cur_level);
  14509. auto [w_ih_value, w_ih_bdim] = unwrapTensorAtLevel(w_ih, cur_level);
  14510. auto [w_hh_value, w_hh_bdim] = unwrapTensorAtLevel(w_hh, cur_level);
  14511. auto [b_ih_value, b_ih_bdim] = unwrapTensorAtLevel(b_ih, cur_level);
  14512. auto [b_hh_value, b_hh_bdim] = unwrapTensorAtLevel(b_hh, cur_level);
  14513. auto [packed_ih_value, packed_ih_bdim] = unwrapTensorAtLevel(packed_ih, cur_level);
  14514. auto [packed_hh_value, packed_hh_bdim] = unwrapTensorAtLevel(packed_hh, cur_level);
  14515. auto [col_offsets_ih_value, col_offsets_ih_bdim] = unwrapTensorAtLevel(col_offsets_ih, cur_level);
  14516. auto [col_offsets_hh_value, col_offsets_hh_bdim] = unwrapTensorAtLevel(col_offsets_hh, cur_level);
  14517. auto results = batch_rule(input_value, input_bdim, hx_value, hx_bdim, w_ih_value, w_ih_bdim, w_hh_value, w_hh_bdim, b_ih_value, b_ih_bdim, b_hh_value, b_hh_bdim, packed_ih_value, packed_ih_bdim, packed_hh_value, packed_hh_bdim, col_offsets_ih_value, col_offsets_ih_bdim, col_offsets_hh_value, col_offsets_hh_bdim, scale_ih, scale_hh, zero_point_ih, zero_point_hh);
  14518. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  14519. }
  14520. template <typename batch_rule_t, batch_rule_t batch_rule>
  14521. ::std::tuple<at::Tensor,at::Tensor> _pack_padded_sequence_generated_plumbing(const at::Tensor & input, const at::Tensor & lengths, bool batch_first) {
  14522. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14523. auto maybe_layer = maybeCurrentDynamicLayer();
  14524. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14525. int64_t cur_level = maybe_layer->layerId();
  14526. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(lengths, cur_level)) {
  14527. return at::_ops::_pack_padded_sequence::call(input, lengths, batch_first);
  14528. }
  14529. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  14530. auto [lengths_value, lengths_bdim] = unwrapTensorAtLevel(lengths, cur_level);
  14531. auto results = batch_rule(input_value, input_bdim, lengths_value, lengths_bdim, batch_first);
  14532. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  14533. }
  14534. template <typename batch_rule_t, batch_rule_t batch_rule>
  14535. at::Tensor _pack_padded_sequence_backward_generated_plumbing(const at::Tensor & grad, c10::SymIntArrayRef input_size, const at::Tensor & batch_sizes, bool batch_first) {
  14536. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14537. auto maybe_layer = maybeCurrentDynamicLayer();
  14538. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14539. int64_t cur_level = maybe_layer->layerId();
  14540. if (!isBatchedAtLevel(grad, cur_level) && !isBatchedAtLevel(batch_sizes, cur_level)) {
  14541. return at::_ops::_pack_padded_sequence_backward::call(grad, input_size, batch_sizes, batch_first);
  14542. }
  14543. auto [grad_value, grad_bdim] = unwrapTensorAtLevel(grad, cur_level);
  14544. auto [batch_sizes_value, batch_sizes_bdim] = unwrapTensorAtLevel(batch_sizes, cur_level);
  14545. auto results = batch_rule(grad_value, grad_bdim, input_size, batch_sizes_value, batch_sizes_bdim, batch_first);
  14546. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  14547. }
  14548. template <typename batch_rule_t, batch_rule_t batch_rule>
  14549. ::std::tuple<at::Tensor,at::Tensor> _pad_packed_sequence_generated_plumbing(const at::Tensor & data, const at::Tensor & batch_sizes, bool batch_first, const at::Scalar & padding_value, int64_t total_length) {
  14550. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14551. auto maybe_layer = maybeCurrentDynamicLayer();
  14552. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14553. int64_t cur_level = maybe_layer->layerId();
  14554. if (!isBatchedAtLevel(data, cur_level) && !isBatchedAtLevel(batch_sizes, cur_level)) {
  14555. return at::_ops::_pad_packed_sequence::call(data, batch_sizes, batch_first, padding_value, total_length);
  14556. }
  14557. auto [data_value, data_bdim] = unwrapTensorAtLevel(data, cur_level);
  14558. auto [batch_sizes_value, batch_sizes_bdim] = unwrapTensorAtLevel(batch_sizes, cur_level);
  14559. auto results = batch_rule(data_value, data_bdim, batch_sizes_value, batch_sizes_bdim, batch_first, padding_value, total_length);
  14560. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  14561. }
  14562. template <typename batch_rule_t, batch_rule_t batch_rule>
  14563. at::Tensor lift_generated_plumbing(const at::Tensor & self) {
  14564. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14565. auto maybe_layer = maybeCurrentDynamicLayer();
  14566. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14567. int64_t cur_level = maybe_layer->layerId();
  14568. if (!isBatchedAtLevel(self, cur_level)) {
  14569. return at::_ops::lift::call(self);
  14570. }
  14571. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  14572. auto results = batch_rule(self_value, self_bdim);
  14573. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  14574. }
  14575. template <typename batch_rule_t, batch_rule_t batch_rule>
  14576. at::Tensor lift_fresh_generated_plumbing(const at::Tensor & self) {
  14577. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14578. auto maybe_layer = maybeCurrentDynamicLayer();
  14579. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14580. int64_t cur_level = maybe_layer->layerId();
  14581. if (!isBatchedAtLevel(self, cur_level)) {
  14582. return at::_ops::lift_fresh::call(self);
  14583. }
  14584. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  14585. auto results = batch_rule(self_value, self_bdim);
  14586. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  14587. }
  14588. template <typename batch_rule_t, batch_rule_t batch_rule>
  14589. at::Tensor lift_fresh_copy_generated_plumbing(const at::Tensor & self) {
  14590. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14591. auto maybe_layer = maybeCurrentDynamicLayer();
  14592. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14593. int64_t cur_level = maybe_layer->layerId();
  14594. if (!isBatchedAtLevel(self, cur_level)) {
  14595. return at::_ops::lift_fresh_copy::call(self);
  14596. }
  14597. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  14598. auto results = batch_rule(self_value, self_bdim);
  14599. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  14600. }
  14601. template <typename batch_rule_t, batch_rule_t batch_rule>
  14602. at::Tensor & masked_fill__Scalar_generated_plumbing(at::Tensor & self, const at::Tensor & mask, const at::Scalar & value) {
  14603. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14604. auto maybe_layer = maybeCurrentDynamicLayer();
  14605. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  14606. int64_t cur_level = maybe_layer->layerId();
  14607. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(mask, cur_level)) {
  14608. return at::_ops::masked_fill__Scalar::call(self, mask, value);
  14609. }
  14610. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  14611. auto [mask_value, mask_bdim] = unwrapTensorAtLevel(mask, cur_level);
  14612. batch_rule(self_value, self_bdim, mask_value, mask_bdim, value);
  14613. return self;
  14614. }
  14615. template <typename batch_rule_t, batch_rule_t batch_rule>
  14616. at::Tensor masked_fill_Scalar_generated_plumbing(const at::Tensor & self, const at::Tensor & mask, const at::Scalar & value) {
  14617. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14618. auto maybe_layer = maybeCurrentDynamicLayer();
  14619. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14620. int64_t cur_level = maybe_layer->layerId();
  14621. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(mask, cur_level)) {
  14622. return at::_ops::masked_fill_Scalar::call(self, mask, value);
  14623. }
  14624. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  14625. auto [mask_value, mask_bdim] = unwrapTensorAtLevel(mask, cur_level);
  14626. auto results = batch_rule(self_value, self_bdim, mask_value, mask_bdim, value);
  14627. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  14628. }
  14629. template <typename batch_rule_t, batch_rule_t batch_rule>
  14630. at::Tensor & masked_fill__Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & mask, const at::Tensor & value) {
  14631. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14632. auto maybe_layer = maybeCurrentDynamicLayer();
  14633. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  14634. int64_t cur_level = maybe_layer->layerId();
  14635. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(mask, cur_level) && !isBatchedAtLevel(value, cur_level)) {
  14636. return at::_ops::masked_fill__Tensor::call(self, mask, value);
  14637. }
  14638. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  14639. auto [mask_value, mask_bdim] = unwrapTensorAtLevel(mask, cur_level);
  14640. auto [value_value, value_bdim] = unwrapTensorAtLevel(value, cur_level);
  14641. batch_rule(self_value, self_bdim, mask_value, mask_bdim, value_value, value_bdim);
  14642. return self;
  14643. }
  14644. template <typename batch_rule_t, batch_rule_t batch_rule>
  14645. at::Tensor masked_fill_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & mask, const at::Tensor & value) {
  14646. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14647. auto maybe_layer = maybeCurrentDynamicLayer();
  14648. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14649. int64_t cur_level = maybe_layer->layerId();
  14650. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(mask, cur_level) && !isBatchedAtLevel(value, cur_level)) {
  14651. return at::_ops::masked_fill_Tensor::call(self, mask, value);
  14652. }
  14653. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  14654. auto [mask_value, mask_bdim] = unwrapTensorAtLevel(mask, cur_level);
  14655. auto [value_value, value_bdim] = unwrapTensorAtLevel(value, cur_level);
  14656. auto results = batch_rule(self_value, self_bdim, mask_value, mask_bdim, value_value, value_bdim);
  14657. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  14658. }
  14659. template <typename batch_rule_t, batch_rule_t batch_rule>
  14660. at::Tensor & masked_scatter__generated_plumbing(at::Tensor & self, const at::Tensor & mask, const at::Tensor & source) {
  14661. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14662. auto maybe_layer = maybeCurrentDynamicLayer();
  14663. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  14664. int64_t cur_level = maybe_layer->layerId();
  14665. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(mask, cur_level) && !isBatchedAtLevel(source, cur_level)) {
  14666. return at::_ops::masked_scatter_::call(self, mask, source);
  14667. }
  14668. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  14669. auto [mask_value, mask_bdim] = unwrapTensorAtLevel(mask, cur_level);
  14670. auto [source_value, source_bdim] = unwrapTensorAtLevel(source, cur_level);
  14671. batch_rule(self_value, self_bdim, mask_value, mask_bdim, source_value, source_bdim);
  14672. return self;
  14673. }
  14674. template <typename batch_rule_t, batch_rule_t batch_rule>
  14675. at::Tensor masked_scatter_generated_plumbing(const at::Tensor & self, const at::Tensor & mask, const at::Tensor & source) {
  14676. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14677. auto maybe_layer = maybeCurrentDynamicLayer();
  14678. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14679. int64_t cur_level = maybe_layer->layerId();
  14680. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(mask, cur_level) && !isBatchedAtLevel(source, cur_level)) {
  14681. return at::_ops::masked_scatter::call(self, mask, source);
  14682. }
  14683. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  14684. auto [mask_value, mask_bdim] = unwrapTensorAtLevel(mask, cur_level);
  14685. auto [source_value, source_bdim] = unwrapTensorAtLevel(source, cur_level);
  14686. auto results = batch_rule(self_value, self_bdim, mask_value, mask_bdim, source_value, source_bdim);
  14687. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  14688. }
  14689. template <typename batch_rule_t, batch_rule_t batch_rule>
  14690. at::Tensor masked_scatter_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & mask, c10::SymIntArrayRef sizes) {
  14691. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14692. auto maybe_layer = maybeCurrentDynamicLayer();
  14693. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14694. int64_t cur_level = maybe_layer->layerId();
  14695. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(mask, cur_level)) {
  14696. return at::_ops::masked_scatter_backward::call(grad_output, mask, sizes);
  14697. }
  14698. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  14699. auto [mask_value, mask_bdim] = unwrapTensorAtLevel(mask, cur_level);
  14700. auto results = batch_rule(grad_output_value, grad_output_bdim, mask_value, mask_bdim, sizes);
  14701. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  14702. }
  14703. template <typename batch_rule_t, batch_rule_t batch_rule>
  14704. at::Tensor _masked_softmax_generated_plumbing(const at::Tensor & self, const at::Tensor & mask, ::std::optional<int64_t> dim, ::std::optional<int64_t> mask_type) {
  14705. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14706. auto maybe_layer = maybeCurrentDynamicLayer();
  14707. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14708. int64_t cur_level = maybe_layer->layerId();
  14709. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(mask, cur_level)) {
  14710. return at::_ops::_masked_softmax::call(self, mask, dim, mask_type);
  14711. }
  14712. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  14713. auto [mask_value, mask_bdim] = unwrapTensorAtLevel(mask, cur_level);
  14714. auto results = batch_rule(self_value, self_bdim, mask_value, mask_bdim, dim, mask_type);
  14715. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  14716. }
  14717. template <typename batch_rule_t, batch_rule_t batch_rule>
  14718. at::Tensor _masked_softmax_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & output, const at::Tensor & mask, ::std::optional<int64_t> dim) {
  14719. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14720. auto maybe_layer = maybeCurrentDynamicLayer();
  14721. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14722. int64_t cur_level = maybe_layer->layerId();
  14723. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(output, cur_level) && !isBatchedAtLevel(mask, cur_level)) {
  14724. return at::_ops::_masked_softmax_backward::call(grad_output, output, mask, dim);
  14725. }
  14726. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  14727. auto [output_value, output_bdim] = unwrapTensorAtLevel(output, cur_level);
  14728. auto [mask_value, mask_bdim] = unwrapTensorAtLevel(mask, cur_level);
  14729. auto results = batch_rule(grad_output_value, grad_output_bdim, output_value, output_bdim, mask_value, mask_bdim, dim);
  14730. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  14731. }
  14732. template <typename batch_rule_t, batch_rule_t batch_rule>
  14733. at::Tensor view_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef size) {
  14734. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14735. auto maybe_layer = maybeCurrentDynamicLayer();
  14736. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14737. int64_t cur_level = maybe_layer->layerId();
  14738. if (!isBatchedAtLevel(self, cur_level)) {
  14739. return at::_ops::view::call(self, size);
  14740. }
  14741. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  14742. auto results = batch_rule(self_value, self_bdim, size);
  14743. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  14744. }
  14745. template <typename batch_rule_t, batch_rule_t batch_rule>
  14746. at::Tensor view_dtype_generated_plumbing(const at::Tensor & self, at::ScalarType dtype) {
  14747. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14748. auto maybe_layer = maybeCurrentDynamicLayer();
  14749. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14750. int64_t cur_level = maybe_layer->layerId();
  14751. if (!isBatchedAtLevel(self, cur_level)) {
  14752. return at::_ops::view_dtype::call(self, dtype);
  14753. }
  14754. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  14755. auto results = batch_rule(self_value, self_bdim, dtype);
  14756. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  14757. }
  14758. template <typename batch_rule_t, batch_rule_t batch_rule>
  14759. at::Tensor & put__generated_plumbing(at::Tensor & self, const at::Tensor & index, const at::Tensor & source, bool accumulate) {
  14760. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14761. auto maybe_layer = maybeCurrentDynamicLayer();
  14762. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  14763. int64_t cur_level = maybe_layer->layerId();
  14764. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(index, cur_level) && !isBatchedAtLevel(source, cur_level)) {
  14765. return at::_ops::put_::call(self, index, source, accumulate);
  14766. }
  14767. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  14768. auto [index_value, index_bdim] = unwrapTensorAtLevel(index, cur_level);
  14769. auto [source_value, source_bdim] = unwrapTensorAtLevel(source, cur_level);
  14770. batch_rule(self_value, self_bdim, index_value, index_bdim, source_value, source_bdim, accumulate);
  14771. return self;
  14772. }
  14773. template <typename batch_rule_t, batch_rule_t batch_rule>
  14774. at::Tensor put_generated_plumbing(const at::Tensor & self, const at::Tensor & index, const at::Tensor & source, bool accumulate) {
  14775. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14776. auto maybe_layer = maybeCurrentDynamicLayer();
  14777. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14778. int64_t cur_level = maybe_layer->layerId();
  14779. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(index, cur_level) && !isBatchedAtLevel(source, cur_level)) {
  14780. return at::_ops::put::call(self, index, source, accumulate);
  14781. }
  14782. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  14783. auto [index_value, index_bdim] = unwrapTensorAtLevel(index, cur_level);
  14784. auto [source_value, source_bdim] = unwrapTensorAtLevel(source, cur_level);
  14785. auto results = batch_rule(self_value, self_bdim, index_value, index_bdim, source_value, source_bdim, accumulate);
  14786. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  14787. }
  14788. template <typename batch_rule_t, batch_rule_t batch_rule>
  14789. at::Tensor & index_add__generated_plumbing(at::Tensor & self, int64_t dim, const at::Tensor & index, const at::Tensor & source, const at::Scalar & alpha) {
  14790. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14791. auto maybe_layer = maybeCurrentDynamicLayer();
  14792. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  14793. int64_t cur_level = maybe_layer->layerId();
  14794. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(index, cur_level) && !isBatchedAtLevel(source, cur_level)) {
  14795. return at::_ops::index_add_::call(self, dim, index, source, alpha);
  14796. }
  14797. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  14798. auto [index_value, index_bdim] = unwrapTensorAtLevel(index, cur_level);
  14799. auto [source_value, source_bdim] = unwrapTensorAtLevel(source, cur_level);
  14800. batch_rule(self_value, self_bdim, dim, index_value, index_bdim, source_value, source_bdim, alpha);
  14801. return self;
  14802. }
  14803. template <typename batch_rule_t, batch_rule_t batch_rule>
  14804. at::Tensor index_add_generated_plumbing(const at::Tensor & self, int64_t dim, const at::Tensor & index, const at::Tensor & source, const at::Scalar & alpha) {
  14805. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14806. auto maybe_layer = maybeCurrentDynamicLayer();
  14807. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14808. int64_t cur_level = maybe_layer->layerId();
  14809. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(index, cur_level) && !isBatchedAtLevel(source, cur_level)) {
  14810. return at::_ops::index_add::call(self, dim, index, source, alpha);
  14811. }
  14812. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  14813. auto [index_value, index_bdim] = unwrapTensorAtLevel(index, cur_level);
  14814. auto [source_value, source_bdim] = unwrapTensorAtLevel(source, cur_level);
  14815. auto results = batch_rule(self_value, self_bdim, dim, index_value, index_bdim, source_value, source_bdim, alpha);
  14816. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  14817. }
  14818. template <typename batch_rule_t, batch_rule_t batch_rule>
  14819. at::Tensor index_add_dimname_generated_plumbing(const at::Tensor & self, at::Dimname dim, const at::Tensor & index, const at::Tensor & source, const at::Scalar & alpha) {
  14820. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14821. auto maybe_layer = maybeCurrentDynamicLayer();
  14822. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14823. int64_t cur_level = maybe_layer->layerId();
  14824. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(index, cur_level) && !isBatchedAtLevel(source, cur_level)) {
  14825. return at::_ops::index_add_dimname::call(self, dim, index, source, alpha);
  14826. }
  14827. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  14828. auto [index_value, index_bdim] = unwrapTensorAtLevel(index, cur_level);
  14829. auto [source_value, source_bdim] = unwrapTensorAtLevel(source, cur_level);
  14830. auto results = batch_rule(self_value, self_bdim, dim, index_value, index_bdim, source_value, source_bdim, alpha);
  14831. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  14832. }
  14833. template <typename batch_rule_t, batch_rule_t batch_rule>
  14834. at::Tensor & index_reduce__generated_plumbing(at::Tensor & self, int64_t dim, const at::Tensor & index, const at::Tensor & source, c10::string_view reduce, bool include_self) {
  14835. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14836. auto maybe_layer = maybeCurrentDynamicLayer();
  14837. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  14838. int64_t cur_level = maybe_layer->layerId();
  14839. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(index, cur_level) && !isBatchedAtLevel(source, cur_level)) {
  14840. return at::_ops::index_reduce_::call(self, dim, index, source, reduce, include_self);
  14841. }
  14842. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  14843. auto [index_value, index_bdim] = unwrapTensorAtLevel(index, cur_level);
  14844. auto [source_value, source_bdim] = unwrapTensorAtLevel(source, cur_level);
  14845. batch_rule(self_value, self_bdim, dim, index_value, index_bdim, source_value, source_bdim, reduce, include_self);
  14846. return self;
  14847. }
  14848. template <typename batch_rule_t, batch_rule_t batch_rule>
  14849. at::Tensor index_reduce_generated_plumbing(const at::Tensor & self, int64_t dim, const at::Tensor & index, const at::Tensor & source, c10::string_view reduce, bool include_self) {
  14850. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14851. auto maybe_layer = maybeCurrentDynamicLayer();
  14852. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14853. int64_t cur_level = maybe_layer->layerId();
  14854. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(index, cur_level) && !isBatchedAtLevel(source, cur_level)) {
  14855. return at::_ops::index_reduce::call(self, dim, index, source, reduce, include_self);
  14856. }
  14857. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  14858. auto [index_value, index_bdim] = unwrapTensorAtLevel(index, cur_level);
  14859. auto [source_value, source_bdim] = unwrapTensorAtLevel(source, cur_level);
  14860. auto results = batch_rule(self_value, self_bdim, dim, index_value, index_bdim, source_value, source_bdim, reduce, include_self);
  14861. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  14862. }
  14863. template <typename batch_rule_t, batch_rule_t batch_rule>
  14864. at::Tensor & index_fill__int_Scalar_generated_plumbing(at::Tensor & self, int64_t dim, const at::Tensor & index, const at::Scalar & value) {
  14865. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14866. auto maybe_layer = maybeCurrentDynamicLayer();
  14867. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  14868. int64_t cur_level = maybe_layer->layerId();
  14869. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(index, cur_level)) {
  14870. return at::_ops::index_fill__int_Scalar::call(self, dim, index, value);
  14871. }
  14872. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  14873. auto [index_value, index_bdim] = unwrapTensorAtLevel(index, cur_level);
  14874. batch_rule(self_value, self_bdim, dim, index_value, index_bdim, value);
  14875. return self;
  14876. }
  14877. template <typename batch_rule_t, batch_rule_t batch_rule>
  14878. at::Tensor index_fill_int_Scalar_generated_plumbing(const at::Tensor & self, int64_t dim, const at::Tensor & index, const at::Scalar & value) {
  14879. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14880. auto maybe_layer = maybeCurrentDynamicLayer();
  14881. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14882. int64_t cur_level = maybe_layer->layerId();
  14883. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(index, cur_level)) {
  14884. return at::_ops::index_fill_int_Scalar::call(self, dim, index, value);
  14885. }
  14886. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  14887. auto [index_value, index_bdim] = unwrapTensorAtLevel(index, cur_level);
  14888. auto results = batch_rule(self_value, self_bdim, dim, index_value, index_bdim, value);
  14889. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  14890. }
  14891. template <typename batch_rule_t, batch_rule_t batch_rule>
  14892. at::Tensor & index_fill__int_Tensor_generated_plumbing(at::Tensor & self, int64_t dim, const at::Tensor & index, const at::Tensor & value) {
  14893. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14894. auto maybe_layer = maybeCurrentDynamicLayer();
  14895. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  14896. int64_t cur_level = maybe_layer->layerId();
  14897. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(index, cur_level) && !isBatchedAtLevel(value, cur_level)) {
  14898. return at::_ops::index_fill__int_Tensor::call(self, dim, index, value);
  14899. }
  14900. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  14901. auto [index_value, index_bdim] = unwrapTensorAtLevel(index, cur_level);
  14902. auto [value_value, value_bdim] = unwrapTensorAtLevel(value, cur_level);
  14903. batch_rule(self_value, self_bdim, dim, index_value, index_bdim, value_value, value_bdim);
  14904. return self;
  14905. }
  14906. template <typename batch_rule_t, batch_rule_t batch_rule>
  14907. at::Tensor index_fill_int_Tensor_generated_plumbing(const at::Tensor & self, int64_t dim, const at::Tensor & index, const at::Tensor & value) {
  14908. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14909. auto maybe_layer = maybeCurrentDynamicLayer();
  14910. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14911. int64_t cur_level = maybe_layer->layerId();
  14912. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(index, cur_level) && !isBatchedAtLevel(value, cur_level)) {
  14913. return at::_ops::index_fill_int_Tensor::call(self, dim, index, value);
  14914. }
  14915. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  14916. auto [index_value, index_bdim] = unwrapTensorAtLevel(index, cur_level);
  14917. auto [value_value, value_bdim] = unwrapTensorAtLevel(value, cur_level);
  14918. auto results = batch_rule(self_value, self_bdim, dim, index_value, index_bdim, value_value, value_bdim);
  14919. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  14920. }
  14921. template <typename batch_rule_t, batch_rule_t batch_rule>
  14922. at::Tensor & index_fill__Dimname_Scalar_generated_plumbing(at::Tensor & self, at::Dimname dim, const at::Tensor & index, const at::Scalar & value) {
  14923. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14924. auto maybe_layer = maybeCurrentDynamicLayer();
  14925. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  14926. int64_t cur_level = maybe_layer->layerId();
  14927. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(index, cur_level)) {
  14928. return at::_ops::index_fill__Dimname_Scalar::call(self, dim, index, value);
  14929. }
  14930. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  14931. auto [index_value, index_bdim] = unwrapTensorAtLevel(index, cur_level);
  14932. batch_rule(self_value, self_bdim, dim, index_value, index_bdim, value);
  14933. return self;
  14934. }
  14935. template <typename batch_rule_t, batch_rule_t batch_rule>
  14936. at::Tensor & index_fill__Dimname_Tensor_generated_plumbing(at::Tensor & self, at::Dimname dim, const at::Tensor & index, const at::Tensor & value) {
  14937. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14938. auto maybe_layer = maybeCurrentDynamicLayer();
  14939. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  14940. int64_t cur_level = maybe_layer->layerId();
  14941. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(index, cur_level) && !isBatchedAtLevel(value, cur_level)) {
  14942. return at::_ops::index_fill__Dimname_Tensor::call(self, dim, index, value);
  14943. }
  14944. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  14945. auto [index_value, index_bdim] = unwrapTensorAtLevel(index, cur_level);
  14946. auto [value_value, value_bdim] = unwrapTensorAtLevel(value, cur_level);
  14947. batch_rule(self_value, self_bdim, dim, index_value, index_bdim, value_value, value_bdim);
  14948. return self;
  14949. }
  14950. template <typename batch_rule_t, batch_rule_t batch_rule>
  14951. at::Tensor index_fill_Dimname_Scalar_generated_plumbing(const at::Tensor & self, at::Dimname dim, const at::Tensor & index, const at::Scalar & value) {
  14952. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14953. auto maybe_layer = maybeCurrentDynamicLayer();
  14954. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14955. int64_t cur_level = maybe_layer->layerId();
  14956. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(index, cur_level)) {
  14957. return at::_ops::index_fill_Dimname_Scalar::call(self, dim, index, value);
  14958. }
  14959. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  14960. auto [index_value, index_bdim] = unwrapTensorAtLevel(index, cur_level);
  14961. auto results = batch_rule(self_value, self_bdim, dim, index_value, index_bdim, value);
  14962. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  14963. }
  14964. template <typename batch_rule_t, batch_rule_t batch_rule>
  14965. at::Tensor index_fill_Dimname_Tensor_generated_plumbing(const at::Tensor & self, at::Dimname dim, const at::Tensor & index, const at::Tensor & value) {
  14966. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14967. auto maybe_layer = maybeCurrentDynamicLayer();
  14968. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14969. int64_t cur_level = maybe_layer->layerId();
  14970. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(index, cur_level) && !isBatchedAtLevel(value, cur_level)) {
  14971. return at::_ops::index_fill_Dimname_Tensor::call(self, dim, index, value);
  14972. }
  14973. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  14974. auto [index_value, index_bdim] = unwrapTensorAtLevel(index, cur_level);
  14975. auto [value_value, value_bdim] = unwrapTensorAtLevel(value, cur_level);
  14976. auto results = batch_rule(self_value, self_bdim, dim, index_value, index_bdim, value_value, value_bdim);
  14977. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  14978. }
  14979. template <typename batch_rule_t, batch_rule_t batch_rule>
  14980. at::Tensor scatter_src_generated_plumbing(const at::Tensor & self, int64_t dim, const at::Tensor & index, const at::Tensor & src) {
  14981. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14982. auto maybe_layer = maybeCurrentDynamicLayer();
  14983. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  14984. int64_t cur_level = maybe_layer->layerId();
  14985. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(index, cur_level) && !isBatchedAtLevel(src, cur_level)) {
  14986. return at::_ops::scatter_src::call(self, dim, index, src);
  14987. }
  14988. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  14989. auto [index_value, index_bdim] = unwrapTensorAtLevel(index, cur_level);
  14990. auto [src_value, src_bdim] = unwrapTensorAtLevel(src, cur_level);
  14991. auto results = batch_rule(self_value, self_bdim, dim, index_value, index_bdim, src_value, src_bdim);
  14992. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  14993. }
  14994. template <typename batch_rule_t, batch_rule_t batch_rule>
  14995. at::Tensor & scatter__src_generated_plumbing(at::Tensor & self, int64_t dim, const at::Tensor & index, const at::Tensor & src) {
  14996. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  14997. auto maybe_layer = maybeCurrentDynamicLayer();
  14998. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  14999. int64_t cur_level = maybe_layer->layerId();
  15000. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(index, cur_level) && !isBatchedAtLevel(src, cur_level)) {
  15001. return at::_ops::scatter__src::call(self, dim, index, src);
  15002. }
  15003. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15004. auto [index_value, index_bdim] = unwrapTensorAtLevel(index, cur_level);
  15005. auto [src_value, src_bdim] = unwrapTensorAtLevel(src, cur_level);
  15006. batch_rule(self_value, self_bdim, dim, index_value, index_bdim, src_value, src_bdim);
  15007. return self;
  15008. }
  15009. template <typename batch_rule_t, batch_rule_t batch_rule>
  15010. at::Tensor scatter_value_generated_plumbing(const at::Tensor & self, int64_t dim, const at::Tensor & index, const at::Scalar & value) {
  15011. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15012. auto maybe_layer = maybeCurrentDynamicLayer();
  15013. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  15014. int64_t cur_level = maybe_layer->layerId();
  15015. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(index, cur_level)) {
  15016. return at::_ops::scatter_value::call(self, dim, index, value);
  15017. }
  15018. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15019. auto [index_value, index_bdim] = unwrapTensorAtLevel(index, cur_level);
  15020. auto results = batch_rule(self_value, self_bdim, dim, index_value, index_bdim, value);
  15021. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  15022. }
  15023. template <typename batch_rule_t, batch_rule_t batch_rule>
  15024. at::Tensor & scatter__value_generated_plumbing(at::Tensor & self, int64_t dim, const at::Tensor & index, const at::Scalar & value) {
  15025. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15026. auto maybe_layer = maybeCurrentDynamicLayer();
  15027. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  15028. int64_t cur_level = maybe_layer->layerId();
  15029. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(index, cur_level)) {
  15030. return at::_ops::scatter__value::call(self, dim, index, value);
  15031. }
  15032. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15033. auto [index_value, index_bdim] = unwrapTensorAtLevel(index, cur_level);
  15034. batch_rule(self_value, self_bdim, dim, index_value, index_bdim, value);
  15035. return self;
  15036. }
  15037. template <typename batch_rule_t, batch_rule_t batch_rule>
  15038. at::Tensor scatter_reduce_generated_plumbing(const at::Tensor & self, int64_t dim, const at::Tensor & index, const at::Tensor & src, c10::string_view reduce) {
  15039. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15040. auto maybe_layer = maybeCurrentDynamicLayer();
  15041. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  15042. int64_t cur_level = maybe_layer->layerId();
  15043. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(index, cur_level) && !isBatchedAtLevel(src, cur_level)) {
  15044. return at::_ops::scatter_reduce::call(self, dim, index, src, reduce);
  15045. }
  15046. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15047. auto [index_value, index_bdim] = unwrapTensorAtLevel(index, cur_level);
  15048. auto [src_value, src_bdim] = unwrapTensorAtLevel(src, cur_level);
  15049. auto results = batch_rule(self_value, self_bdim, dim, index_value, index_bdim, src_value, src_bdim, reduce);
  15050. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  15051. }
  15052. template <typename batch_rule_t, batch_rule_t batch_rule>
  15053. at::Tensor & scatter__reduce_generated_plumbing(at::Tensor & self, int64_t dim, const at::Tensor & index, const at::Tensor & src, c10::string_view reduce) {
  15054. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15055. auto maybe_layer = maybeCurrentDynamicLayer();
  15056. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  15057. int64_t cur_level = maybe_layer->layerId();
  15058. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(index, cur_level) && !isBatchedAtLevel(src, cur_level)) {
  15059. return at::_ops::scatter__reduce::call(self, dim, index, src, reduce);
  15060. }
  15061. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15062. auto [index_value, index_bdim] = unwrapTensorAtLevel(index, cur_level);
  15063. auto [src_value, src_bdim] = unwrapTensorAtLevel(src, cur_level);
  15064. batch_rule(self_value, self_bdim, dim, index_value, index_bdim, src_value, src_bdim, reduce);
  15065. return self;
  15066. }
  15067. template <typename batch_rule_t, batch_rule_t batch_rule>
  15068. at::Tensor scatter_value_reduce_generated_plumbing(const at::Tensor & self, int64_t dim, const at::Tensor & index, const at::Scalar & value, c10::string_view reduce) {
  15069. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15070. auto maybe_layer = maybeCurrentDynamicLayer();
  15071. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  15072. int64_t cur_level = maybe_layer->layerId();
  15073. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(index, cur_level)) {
  15074. return at::_ops::scatter_value_reduce::call(self, dim, index, value, reduce);
  15075. }
  15076. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15077. auto [index_value, index_bdim] = unwrapTensorAtLevel(index, cur_level);
  15078. auto results = batch_rule(self_value, self_bdim, dim, index_value, index_bdim, value, reduce);
  15079. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  15080. }
  15081. template <typename batch_rule_t, batch_rule_t batch_rule>
  15082. at::Tensor & scatter__value_reduce_generated_plumbing(at::Tensor & self, int64_t dim, const at::Tensor & index, const at::Scalar & value, c10::string_view reduce) {
  15083. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15084. auto maybe_layer = maybeCurrentDynamicLayer();
  15085. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  15086. int64_t cur_level = maybe_layer->layerId();
  15087. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(index, cur_level)) {
  15088. return at::_ops::scatter__value_reduce::call(self, dim, index, value, reduce);
  15089. }
  15090. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15091. auto [index_value, index_bdim] = unwrapTensorAtLevel(index, cur_level);
  15092. batch_rule(self_value, self_bdim, dim, index_value, index_bdim, value, reduce);
  15093. return self;
  15094. }
  15095. template <typename batch_rule_t, batch_rule_t batch_rule>
  15096. at::Tensor scatter_dimname_src_generated_plumbing(const at::Tensor & self, at::Dimname dim, const at::Tensor & index, const at::Tensor & src) {
  15097. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15098. auto maybe_layer = maybeCurrentDynamicLayer();
  15099. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  15100. int64_t cur_level = maybe_layer->layerId();
  15101. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(index, cur_level) && !isBatchedAtLevel(src, cur_level)) {
  15102. return at::_ops::scatter_dimname_src::call(self, dim, index, src);
  15103. }
  15104. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15105. auto [index_value, index_bdim] = unwrapTensorAtLevel(index, cur_level);
  15106. auto [src_value, src_bdim] = unwrapTensorAtLevel(src, cur_level);
  15107. auto results = batch_rule(self_value, self_bdim, dim, index_value, index_bdim, src_value, src_bdim);
  15108. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  15109. }
  15110. template <typename batch_rule_t, batch_rule_t batch_rule>
  15111. at::Tensor scatter_dimname_value_generated_plumbing(const at::Tensor & self, at::Dimname dim, const at::Tensor & index, const at::Scalar & value) {
  15112. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15113. auto maybe_layer = maybeCurrentDynamicLayer();
  15114. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  15115. int64_t cur_level = maybe_layer->layerId();
  15116. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(index, cur_level)) {
  15117. return at::_ops::scatter_dimname_value::call(self, dim, index, value);
  15118. }
  15119. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15120. auto [index_value, index_bdim] = unwrapTensorAtLevel(index, cur_level);
  15121. auto results = batch_rule(self_value, self_bdim, dim, index_value, index_bdim, value);
  15122. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  15123. }
  15124. template <typename batch_rule_t, batch_rule_t batch_rule>
  15125. at::Tensor scatter_add_generated_plumbing(const at::Tensor & self, int64_t dim, const at::Tensor & index, const at::Tensor & src) {
  15126. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15127. auto maybe_layer = maybeCurrentDynamicLayer();
  15128. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  15129. int64_t cur_level = maybe_layer->layerId();
  15130. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(index, cur_level) && !isBatchedAtLevel(src, cur_level)) {
  15131. return at::_ops::scatter_add::call(self, dim, index, src);
  15132. }
  15133. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15134. auto [index_value, index_bdim] = unwrapTensorAtLevel(index, cur_level);
  15135. auto [src_value, src_bdim] = unwrapTensorAtLevel(src, cur_level);
  15136. auto results = batch_rule(self_value, self_bdim, dim, index_value, index_bdim, src_value, src_bdim);
  15137. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  15138. }
  15139. template <typename batch_rule_t, batch_rule_t batch_rule>
  15140. at::Tensor & scatter_add__generated_plumbing(at::Tensor & self, int64_t dim, const at::Tensor & index, const at::Tensor & src) {
  15141. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15142. auto maybe_layer = maybeCurrentDynamicLayer();
  15143. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  15144. int64_t cur_level = maybe_layer->layerId();
  15145. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(index, cur_level) && !isBatchedAtLevel(src, cur_level)) {
  15146. return at::_ops::scatter_add_::call(self, dim, index, src);
  15147. }
  15148. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15149. auto [index_value, index_bdim] = unwrapTensorAtLevel(index, cur_level);
  15150. auto [src_value, src_bdim] = unwrapTensorAtLevel(src, cur_level);
  15151. batch_rule(self_value, self_bdim, dim, index_value, index_bdim, src_value, src_bdim);
  15152. return self;
  15153. }
  15154. template <typename batch_rule_t, batch_rule_t batch_rule>
  15155. at::Tensor scatter_add_dimname_generated_plumbing(const at::Tensor & self, at::Dimname dim, const at::Tensor & index, const at::Tensor & src) {
  15156. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15157. auto maybe_layer = maybeCurrentDynamicLayer();
  15158. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  15159. int64_t cur_level = maybe_layer->layerId();
  15160. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(index, cur_level) && !isBatchedAtLevel(src, cur_level)) {
  15161. return at::_ops::scatter_add_dimname::call(self, dim, index, src);
  15162. }
  15163. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15164. auto [index_value, index_bdim] = unwrapTensorAtLevel(index, cur_level);
  15165. auto [src_value, src_bdim] = unwrapTensorAtLevel(src, cur_level);
  15166. auto results = batch_rule(self_value, self_bdim, dim, index_value, index_bdim, src_value, src_bdim);
  15167. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  15168. }
  15169. template <typename batch_rule_t, batch_rule_t batch_rule>
  15170. at::Tensor scatter_reduce_two_generated_plumbing(const at::Tensor & self, int64_t dim, const at::Tensor & index, const at::Tensor & src, c10::string_view reduce, bool include_self) {
  15171. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15172. auto maybe_layer = maybeCurrentDynamicLayer();
  15173. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  15174. int64_t cur_level = maybe_layer->layerId();
  15175. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(index, cur_level) && !isBatchedAtLevel(src, cur_level)) {
  15176. return at::_ops::scatter_reduce_two::call(self, dim, index, src, reduce, include_self);
  15177. }
  15178. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15179. auto [index_value, index_bdim] = unwrapTensorAtLevel(index, cur_level);
  15180. auto [src_value, src_bdim] = unwrapTensorAtLevel(src, cur_level);
  15181. auto results = batch_rule(self_value, self_bdim, dim, index_value, index_bdim, src_value, src_bdim, reduce, include_self);
  15182. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  15183. }
  15184. template <typename batch_rule_t, batch_rule_t batch_rule>
  15185. at::Tensor & scatter_reduce__two_generated_plumbing(at::Tensor & self, int64_t dim, const at::Tensor & index, const at::Tensor & src, c10::string_view reduce, bool include_self) {
  15186. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15187. auto maybe_layer = maybeCurrentDynamicLayer();
  15188. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  15189. int64_t cur_level = maybe_layer->layerId();
  15190. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(index, cur_level) && !isBatchedAtLevel(src, cur_level)) {
  15191. return at::_ops::scatter_reduce__two::call(self, dim, index, src, reduce, include_self);
  15192. }
  15193. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15194. auto [index_value, index_bdim] = unwrapTensorAtLevel(index, cur_level);
  15195. auto [src_value, src_bdim] = unwrapTensorAtLevel(src, cur_level);
  15196. batch_rule(self_value, self_bdim, dim, index_value, index_bdim, src_value, src_bdim, reduce, include_self);
  15197. return self;
  15198. }
  15199. template <typename batch_rule_t, batch_rule_t batch_rule>
  15200. at::Tensor & eq__Scalar_generated_plumbing(at::Tensor & self, const at::Scalar & other) {
  15201. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15202. auto maybe_layer = maybeCurrentDynamicLayer();
  15203. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  15204. int64_t cur_level = maybe_layer->layerId();
  15205. if (!isBatchedAtLevel(self, cur_level)) {
  15206. return at::_ops::eq__Scalar::call(self, other);
  15207. }
  15208. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15209. batch_rule(self_value, self_bdim, other);
  15210. return self;
  15211. }
  15212. template <typename batch_rule_t, batch_rule_t batch_rule>
  15213. at::Tensor & eq__Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  15214. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15215. auto maybe_layer = maybeCurrentDynamicLayer();
  15216. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  15217. int64_t cur_level = maybe_layer->layerId();
  15218. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  15219. return at::_ops::eq__Tensor::call(self, other);
  15220. }
  15221. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15222. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  15223. batch_rule(self_value, self_bdim, other_value, other_bdim);
  15224. return self;
  15225. }
  15226. template <typename batch_rule_t, batch_rule_t batch_rule>
  15227. at::Tensor bitwise_and_Scalar_generated_plumbing(const at::Tensor & self, const at::Scalar & other) {
  15228. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15229. auto maybe_layer = maybeCurrentDynamicLayer();
  15230. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  15231. int64_t cur_level = maybe_layer->layerId();
  15232. if (!isBatchedAtLevel(self, cur_level)) {
  15233. return at::_ops::bitwise_and_Scalar::call(self, other);
  15234. }
  15235. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15236. auto results = batch_rule(self_value, self_bdim, other);
  15237. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  15238. }
  15239. template <typename batch_rule_t, batch_rule_t batch_rule>
  15240. at::Tensor bitwise_and_Scalar_Tensor_generated_plumbing(const at::Scalar & self, const at::Tensor & other) {
  15241. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15242. auto maybe_layer = maybeCurrentDynamicLayer();
  15243. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  15244. int64_t cur_level = maybe_layer->layerId();
  15245. if (!isBatchedAtLevel(other, cur_level)) {
  15246. return at::_ops::bitwise_and_Scalar_Tensor::call(self, other);
  15247. }
  15248. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  15249. auto results = batch_rule(self, other_value, other_bdim);
  15250. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  15251. }
  15252. template <typename batch_rule_t, batch_rule_t batch_rule>
  15253. at::Tensor bitwise_and_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  15254. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15255. auto maybe_layer = maybeCurrentDynamicLayer();
  15256. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  15257. int64_t cur_level = maybe_layer->layerId();
  15258. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  15259. return at::_ops::bitwise_and_Tensor::call(self, other);
  15260. }
  15261. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15262. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  15263. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  15264. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  15265. }
  15266. template <typename batch_rule_t, batch_rule_t batch_rule>
  15267. at::Tensor & bitwise_and__Scalar_generated_plumbing(at::Tensor & self, const at::Scalar & other) {
  15268. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15269. auto maybe_layer = maybeCurrentDynamicLayer();
  15270. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  15271. int64_t cur_level = maybe_layer->layerId();
  15272. if (!isBatchedAtLevel(self, cur_level)) {
  15273. return at::_ops::bitwise_and__Scalar::call(self, other);
  15274. }
  15275. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15276. batch_rule(self_value, self_bdim, other);
  15277. return self;
  15278. }
  15279. template <typename batch_rule_t, batch_rule_t batch_rule>
  15280. at::Tensor & bitwise_and__Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  15281. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15282. auto maybe_layer = maybeCurrentDynamicLayer();
  15283. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  15284. int64_t cur_level = maybe_layer->layerId();
  15285. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  15286. return at::_ops::bitwise_and__Tensor::call(self, other);
  15287. }
  15288. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15289. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  15290. batch_rule(self_value, self_bdim, other_value, other_bdim);
  15291. return self;
  15292. }
  15293. template <typename batch_rule_t, batch_rule_t batch_rule>
  15294. at::Tensor __and___Scalar_generated_plumbing(const at::Tensor & self, const at::Scalar & other) {
  15295. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15296. auto maybe_layer = maybeCurrentDynamicLayer();
  15297. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  15298. int64_t cur_level = maybe_layer->layerId();
  15299. if (!isBatchedAtLevel(self, cur_level)) {
  15300. return at::_ops::__and___Scalar::call(self, other);
  15301. }
  15302. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15303. auto results = batch_rule(self_value, self_bdim, other);
  15304. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  15305. }
  15306. template <typename batch_rule_t, batch_rule_t batch_rule>
  15307. at::Tensor __and___Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  15308. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15309. auto maybe_layer = maybeCurrentDynamicLayer();
  15310. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  15311. int64_t cur_level = maybe_layer->layerId();
  15312. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  15313. return at::_ops::__and___Tensor::call(self, other);
  15314. }
  15315. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15316. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  15317. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  15318. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  15319. }
  15320. template <typename batch_rule_t, batch_rule_t batch_rule>
  15321. at::Tensor & __iand___Scalar_generated_plumbing(at::Tensor & self, const at::Scalar & other) {
  15322. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15323. auto maybe_layer = maybeCurrentDynamicLayer();
  15324. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  15325. int64_t cur_level = maybe_layer->layerId();
  15326. if (!isBatchedAtLevel(self, cur_level)) {
  15327. return at::_ops::__iand___Scalar::call(self, other);
  15328. }
  15329. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15330. batch_rule(self_value, self_bdim, other);
  15331. return self;
  15332. }
  15333. template <typename batch_rule_t, batch_rule_t batch_rule>
  15334. at::Tensor & __iand___Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  15335. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15336. auto maybe_layer = maybeCurrentDynamicLayer();
  15337. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  15338. int64_t cur_level = maybe_layer->layerId();
  15339. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  15340. return at::_ops::__iand___Tensor::call(self, other);
  15341. }
  15342. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15343. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  15344. batch_rule(self_value, self_bdim, other_value, other_bdim);
  15345. return self;
  15346. }
  15347. template <typename batch_rule_t, batch_rule_t batch_rule>
  15348. at::Tensor bitwise_or_Scalar_generated_plumbing(const at::Tensor & self, const at::Scalar & other) {
  15349. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15350. auto maybe_layer = maybeCurrentDynamicLayer();
  15351. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  15352. int64_t cur_level = maybe_layer->layerId();
  15353. if (!isBatchedAtLevel(self, cur_level)) {
  15354. return at::_ops::bitwise_or_Scalar::call(self, other);
  15355. }
  15356. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15357. auto results = batch_rule(self_value, self_bdim, other);
  15358. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  15359. }
  15360. template <typename batch_rule_t, batch_rule_t batch_rule>
  15361. at::Tensor bitwise_or_Scalar_Tensor_generated_plumbing(const at::Scalar & self, const at::Tensor & other) {
  15362. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15363. auto maybe_layer = maybeCurrentDynamicLayer();
  15364. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  15365. int64_t cur_level = maybe_layer->layerId();
  15366. if (!isBatchedAtLevel(other, cur_level)) {
  15367. return at::_ops::bitwise_or_Scalar_Tensor::call(self, other);
  15368. }
  15369. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  15370. auto results = batch_rule(self, other_value, other_bdim);
  15371. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  15372. }
  15373. template <typename batch_rule_t, batch_rule_t batch_rule>
  15374. at::Tensor bitwise_or_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  15375. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15376. auto maybe_layer = maybeCurrentDynamicLayer();
  15377. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  15378. int64_t cur_level = maybe_layer->layerId();
  15379. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  15380. return at::_ops::bitwise_or_Tensor::call(self, other);
  15381. }
  15382. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15383. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  15384. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  15385. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  15386. }
  15387. template <typename batch_rule_t, batch_rule_t batch_rule>
  15388. at::Tensor & bitwise_or__Scalar_generated_plumbing(at::Tensor & self, const at::Scalar & other) {
  15389. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15390. auto maybe_layer = maybeCurrentDynamicLayer();
  15391. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  15392. int64_t cur_level = maybe_layer->layerId();
  15393. if (!isBatchedAtLevel(self, cur_level)) {
  15394. return at::_ops::bitwise_or__Scalar::call(self, other);
  15395. }
  15396. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15397. batch_rule(self_value, self_bdim, other);
  15398. return self;
  15399. }
  15400. template <typename batch_rule_t, batch_rule_t batch_rule>
  15401. at::Tensor & bitwise_or__Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  15402. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15403. auto maybe_layer = maybeCurrentDynamicLayer();
  15404. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  15405. int64_t cur_level = maybe_layer->layerId();
  15406. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  15407. return at::_ops::bitwise_or__Tensor::call(self, other);
  15408. }
  15409. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15410. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  15411. batch_rule(self_value, self_bdim, other_value, other_bdim);
  15412. return self;
  15413. }
  15414. template <typename batch_rule_t, batch_rule_t batch_rule>
  15415. at::Tensor __or___Scalar_generated_plumbing(const at::Tensor & self, const at::Scalar & other) {
  15416. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15417. auto maybe_layer = maybeCurrentDynamicLayer();
  15418. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  15419. int64_t cur_level = maybe_layer->layerId();
  15420. if (!isBatchedAtLevel(self, cur_level)) {
  15421. return at::_ops::__or___Scalar::call(self, other);
  15422. }
  15423. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15424. auto results = batch_rule(self_value, self_bdim, other);
  15425. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  15426. }
  15427. template <typename batch_rule_t, batch_rule_t batch_rule>
  15428. at::Tensor __or___Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  15429. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15430. auto maybe_layer = maybeCurrentDynamicLayer();
  15431. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  15432. int64_t cur_level = maybe_layer->layerId();
  15433. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  15434. return at::_ops::__or___Tensor::call(self, other);
  15435. }
  15436. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15437. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  15438. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  15439. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  15440. }
  15441. template <typename batch_rule_t, batch_rule_t batch_rule>
  15442. at::Tensor & __ior___Scalar_generated_plumbing(at::Tensor & self, const at::Scalar & other) {
  15443. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15444. auto maybe_layer = maybeCurrentDynamicLayer();
  15445. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  15446. int64_t cur_level = maybe_layer->layerId();
  15447. if (!isBatchedAtLevel(self, cur_level)) {
  15448. return at::_ops::__ior___Scalar::call(self, other);
  15449. }
  15450. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15451. batch_rule(self_value, self_bdim, other);
  15452. return self;
  15453. }
  15454. template <typename batch_rule_t, batch_rule_t batch_rule>
  15455. at::Tensor & __ior___Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  15456. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15457. auto maybe_layer = maybeCurrentDynamicLayer();
  15458. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  15459. int64_t cur_level = maybe_layer->layerId();
  15460. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  15461. return at::_ops::__ior___Tensor::call(self, other);
  15462. }
  15463. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15464. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  15465. batch_rule(self_value, self_bdim, other_value, other_bdim);
  15466. return self;
  15467. }
  15468. template <typename batch_rule_t, batch_rule_t batch_rule>
  15469. at::Tensor bitwise_xor_Scalar_generated_plumbing(const at::Tensor & self, const at::Scalar & other) {
  15470. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15471. auto maybe_layer = maybeCurrentDynamicLayer();
  15472. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  15473. int64_t cur_level = maybe_layer->layerId();
  15474. if (!isBatchedAtLevel(self, cur_level)) {
  15475. return at::_ops::bitwise_xor_Scalar::call(self, other);
  15476. }
  15477. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15478. auto results = batch_rule(self_value, self_bdim, other);
  15479. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  15480. }
  15481. template <typename batch_rule_t, batch_rule_t batch_rule>
  15482. at::Tensor bitwise_xor_Scalar_Tensor_generated_plumbing(const at::Scalar & self, const at::Tensor & other) {
  15483. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15484. auto maybe_layer = maybeCurrentDynamicLayer();
  15485. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  15486. int64_t cur_level = maybe_layer->layerId();
  15487. if (!isBatchedAtLevel(other, cur_level)) {
  15488. return at::_ops::bitwise_xor_Scalar_Tensor::call(self, other);
  15489. }
  15490. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  15491. auto results = batch_rule(self, other_value, other_bdim);
  15492. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  15493. }
  15494. template <typename batch_rule_t, batch_rule_t batch_rule>
  15495. at::Tensor bitwise_xor_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  15496. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15497. auto maybe_layer = maybeCurrentDynamicLayer();
  15498. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  15499. int64_t cur_level = maybe_layer->layerId();
  15500. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  15501. return at::_ops::bitwise_xor_Tensor::call(self, other);
  15502. }
  15503. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15504. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  15505. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  15506. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  15507. }
  15508. template <typename batch_rule_t, batch_rule_t batch_rule>
  15509. at::Tensor & bitwise_xor__Scalar_generated_plumbing(at::Tensor & self, const at::Scalar & other) {
  15510. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15511. auto maybe_layer = maybeCurrentDynamicLayer();
  15512. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  15513. int64_t cur_level = maybe_layer->layerId();
  15514. if (!isBatchedAtLevel(self, cur_level)) {
  15515. return at::_ops::bitwise_xor__Scalar::call(self, other);
  15516. }
  15517. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15518. batch_rule(self_value, self_bdim, other);
  15519. return self;
  15520. }
  15521. template <typename batch_rule_t, batch_rule_t batch_rule>
  15522. at::Tensor & bitwise_xor__Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  15523. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15524. auto maybe_layer = maybeCurrentDynamicLayer();
  15525. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  15526. int64_t cur_level = maybe_layer->layerId();
  15527. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  15528. return at::_ops::bitwise_xor__Tensor::call(self, other);
  15529. }
  15530. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15531. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  15532. batch_rule(self_value, self_bdim, other_value, other_bdim);
  15533. return self;
  15534. }
  15535. template <typename batch_rule_t, batch_rule_t batch_rule>
  15536. at::Tensor __xor___Scalar_generated_plumbing(const at::Tensor & self, const at::Scalar & other) {
  15537. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15538. auto maybe_layer = maybeCurrentDynamicLayer();
  15539. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  15540. int64_t cur_level = maybe_layer->layerId();
  15541. if (!isBatchedAtLevel(self, cur_level)) {
  15542. return at::_ops::__xor___Scalar::call(self, other);
  15543. }
  15544. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15545. auto results = batch_rule(self_value, self_bdim, other);
  15546. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  15547. }
  15548. template <typename batch_rule_t, batch_rule_t batch_rule>
  15549. at::Tensor __xor___Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  15550. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15551. auto maybe_layer = maybeCurrentDynamicLayer();
  15552. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  15553. int64_t cur_level = maybe_layer->layerId();
  15554. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  15555. return at::_ops::__xor___Tensor::call(self, other);
  15556. }
  15557. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15558. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  15559. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  15560. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  15561. }
  15562. template <typename batch_rule_t, batch_rule_t batch_rule>
  15563. at::Tensor & __ixor___Scalar_generated_plumbing(at::Tensor & self, const at::Scalar & other) {
  15564. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15565. auto maybe_layer = maybeCurrentDynamicLayer();
  15566. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  15567. int64_t cur_level = maybe_layer->layerId();
  15568. if (!isBatchedAtLevel(self, cur_level)) {
  15569. return at::_ops::__ixor___Scalar::call(self, other);
  15570. }
  15571. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15572. batch_rule(self_value, self_bdim, other);
  15573. return self;
  15574. }
  15575. template <typename batch_rule_t, batch_rule_t batch_rule>
  15576. at::Tensor & __ixor___Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  15577. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15578. auto maybe_layer = maybeCurrentDynamicLayer();
  15579. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  15580. int64_t cur_level = maybe_layer->layerId();
  15581. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  15582. return at::_ops::__ixor___Tensor::call(self, other);
  15583. }
  15584. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15585. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  15586. batch_rule(self_value, self_bdim, other_value, other_bdim);
  15587. return self;
  15588. }
  15589. template <typename batch_rule_t, batch_rule_t batch_rule>
  15590. at::Tensor __lshift___Scalar_generated_plumbing(const at::Tensor & self, const at::Scalar & other) {
  15591. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15592. auto maybe_layer = maybeCurrentDynamicLayer();
  15593. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  15594. int64_t cur_level = maybe_layer->layerId();
  15595. if (!isBatchedAtLevel(self, cur_level)) {
  15596. return at::_ops::__lshift___Scalar::call(self, other);
  15597. }
  15598. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15599. auto results = batch_rule(self_value, self_bdim, other);
  15600. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  15601. }
  15602. template <typename batch_rule_t, batch_rule_t batch_rule>
  15603. at::Tensor __lshift___Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  15604. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15605. auto maybe_layer = maybeCurrentDynamicLayer();
  15606. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  15607. int64_t cur_level = maybe_layer->layerId();
  15608. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  15609. return at::_ops::__lshift___Tensor::call(self, other);
  15610. }
  15611. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15612. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  15613. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  15614. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  15615. }
  15616. template <typename batch_rule_t, batch_rule_t batch_rule>
  15617. at::Tensor & __ilshift___Scalar_generated_plumbing(at::Tensor & self, const at::Scalar & other) {
  15618. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15619. auto maybe_layer = maybeCurrentDynamicLayer();
  15620. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  15621. int64_t cur_level = maybe_layer->layerId();
  15622. if (!isBatchedAtLevel(self, cur_level)) {
  15623. return at::_ops::__ilshift___Scalar::call(self, other);
  15624. }
  15625. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15626. batch_rule(self_value, self_bdim, other);
  15627. return self;
  15628. }
  15629. template <typename batch_rule_t, batch_rule_t batch_rule>
  15630. at::Tensor & __ilshift___Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  15631. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15632. auto maybe_layer = maybeCurrentDynamicLayer();
  15633. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  15634. int64_t cur_level = maybe_layer->layerId();
  15635. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  15636. return at::_ops::__ilshift___Tensor::call(self, other);
  15637. }
  15638. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15639. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  15640. batch_rule(self_value, self_bdim, other_value, other_bdim);
  15641. return self;
  15642. }
  15643. template <typename batch_rule_t, batch_rule_t batch_rule>
  15644. at::Tensor bitwise_left_shift_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  15645. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15646. auto maybe_layer = maybeCurrentDynamicLayer();
  15647. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  15648. int64_t cur_level = maybe_layer->layerId();
  15649. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  15650. return at::_ops::bitwise_left_shift_Tensor::call(self, other);
  15651. }
  15652. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15653. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  15654. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  15655. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  15656. }
  15657. template <typename batch_rule_t, batch_rule_t batch_rule>
  15658. at::Tensor & bitwise_left_shift__Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  15659. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15660. auto maybe_layer = maybeCurrentDynamicLayer();
  15661. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  15662. int64_t cur_level = maybe_layer->layerId();
  15663. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  15664. return at::_ops::bitwise_left_shift__Tensor::call(self, other);
  15665. }
  15666. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15667. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  15668. batch_rule(self_value, self_bdim, other_value, other_bdim);
  15669. return self;
  15670. }
  15671. template <typename batch_rule_t, batch_rule_t batch_rule>
  15672. at::Tensor bitwise_left_shift_Tensor_Scalar_generated_plumbing(const at::Tensor & self, const at::Scalar & other) {
  15673. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15674. auto maybe_layer = maybeCurrentDynamicLayer();
  15675. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  15676. int64_t cur_level = maybe_layer->layerId();
  15677. if (!isBatchedAtLevel(self, cur_level)) {
  15678. return at::_ops::bitwise_left_shift_Tensor_Scalar::call(self, other);
  15679. }
  15680. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15681. auto results = batch_rule(self_value, self_bdim, other);
  15682. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  15683. }
  15684. template <typename batch_rule_t, batch_rule_t batch_rule>
  15685. at::Tensor & bitwise_left_shift__Tensor_Scalar_generated_plumbing(at::Tensor & self, const at::Scalar & other) {
  15686. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15687. auto maybe_layer = maybeCurrentDynamicLayer();
  15688. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  15689. int64_t cur_level = maybe_layer->layerId();
  15690. if (!isBatchedAtLevel(self, cur_level)) {
  15691. return at::_ops::bitwise_left_shift__Tensor_Scalar::call(self, other);
  15692. }
  15693. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15694. batch_rule(self_value, self_bdim, other);
  15695. return self;
  15696. }
  15697. template <typename batch_rule_t, batch_rule_t batch_rule>
  15698. at::Tensor bitwise_left_shift_Scalar_Tensor_generated_plumbing(const at::Scalar & self, const at::Tensor & other) {
  15699. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15700. auto maybe_layer = maybeCurrentDynamicLayer();
  15701. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  15702. int64_t cur_level = maybe_layer->layerId();
  15703. if (!isBatchedAtLevel(other, cur_level)) {
  15704. return at::_ops::bitwise_left_shift_Scalar_Tensor::call(self, other);
  15705. }
  15706. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  15707. auto results = batch_rule(self, other_value, other_bdim);
  15708. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  15709. }
  15710. template <typename batch_rule_t, batch_rule_t batch_rule>
  15711. at::Tensor __rshift___Scalar_generated_plumbing(const at::Tensor & self, const at::Scalar & other) {
  15712. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15713. auto maybe_layer = maybeCurrentDynamicLayer();
  15714. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  15715. int64_t cur_level = maybe_layer->layerId();
  15716. if (!isBatchedAtLevel(self, cur_level)) {
  15717. return at::_ops::__rshift___Scalar::call(self, other);
  15718. }
  15719. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15720. auto results = batch_rule(self_value, self_bdim, other);
  15721. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  15722. }
  15723. template <typename batch_rule_t, batch_rule_t batch_rule>
  15724. at::Tensor __rshift___Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  15725. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15726. auto maybe_layer = maybeCurrentDynamicLayer();
  15727. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  15728. int64_t cur_level = maybe_layer->layerId();
  15729. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  15730. return at::_ops::__rshift___Tensor::call(self, other);
  15731. }
  15732. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15733. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  15734. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  15735. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  15736. }
  15737. template <typename batch_rule_t, batch_rule_t batch_rule>
  15738. at::Tensor & __irshift___Scalar_generated_plumbing(at::Tensor & self, const at::Scalar & other) {
  15739. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15740. auto maybe_layer = maybeCurrentDynamicLayer();
  15741. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  15742. int64_t cur_level = maybe_layer->layerId();
  15743. if (!isBatchedAtLevel(self, cur_level)) {
  15744. return at::_ops::__irshift___Scalar::call(self, other);
  15745. }
  15746. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15747. batch_rule(self_value, self_bdim, other);
  15748. return self;
  15749. }
  15750. template <typename batch_rule_t, batch_rule_t batch_rule>
  15751. at::Tensor & __irshift___Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  15752. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15753. auto maybe_layer = maybeCurrentDynamicLayer();
  15754. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  15755. int64_t cur_level = maybe_layer->layerId();
  15756. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  15757. return at::_ops::__irshift___Tensor::call(self, other);
  15758. }
  15759. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15760. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  15761. batch_rule(self_value, self_bdim, other_value, other_bdim);
  15762. return self;
  15763. }
  15764. template <typename batch_rule_t, batch_rule_t batch_rule>
  15765. at::Tensor bitwise_right_shift_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  15766. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15767. auto maybe_layer = maybeCurrentDynamicLayer();
  15768. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  15769. int64_t cur_level = maybe_layer->layerId();
  15770. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  15771. return at::_ops::bitwise_right_shift_Tensor::call(self, other);
  15772. }
  15773. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15774. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  15775. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  15776. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  15777. }
  15778. template <typename batch_rule_t, batch_rule_t batch_rule>
  15779. at::Tensor & bitwise_right_shift__Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  15780. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15781. auto maybe_layer = maybeCurrentDynamicLayer();
  15782. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  15783. int64_t cur_level = maybe_layer->layerId();
  15784. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  15785. return at::_ops::bitwise_right_shift__Tensor::call(self, other);
  15786. }
  15787. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15788. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  15789. batch_rule(self_value, self_bdim, other_value, other_bdim);
  15790. return self;
  15791. }
  15792. template <typename batch_rule_t, batch_rule_t batch_rule>
  15793. at::Tensor bitwise_right_shift_Tensor_Scalar_generated_plumbing(const at::Tensor & self, const at::Scalar & other) {
  15794. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15795. auto maybe_layer = maybeCurrentDynamicLayer();
  15796. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  15797. int64_t cur_level = maybe_layer->layerId();
  15798. if (!isBatchedAtLevel(self, cur_level)) {
  15799. return at::_ops::bitwise_right_shift_Tensor_Scalar::call(self, other);
  15800. }
  15801. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15802. auto results = batch_rule(self_value, self_bdim, other);
  15803. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  15804. }
  15805. template <typename batch_rule_t, batch_rule_t batch_rule>
  15806. at::Tensor & bitwise_right_shift__Tensor_Scalar_generated_plumbing(at::Tensor & self, const at::Scalar & other) {
  15807. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15808. auto maybe_layer = maybeCurrentDynamicLayer();
  15809. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  15810. int64_t cur_level = maybe_layer->layerId();
  15811. if (!isBatchedAtLevel(self, cur_level)) {
  15812. return at::_ops::bitwise_right_shift__Tensor_Scalar::call(self, other);
  15813. }
  15814. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15815. batch_rule(self_value, self_bdim, other);
  15816. return self;
  15817. }
  15818. template <typename batch_rule_t, batch_rule_t batch_rule>
  15819. at::Tensor bitwise_right_shift_Scalar_Tensor_generated_plumbing(const at::Scalar & self, const at::Tensor & other) {
  15820. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15821. auto maybe_layer = maybeCurrentDynamicLayer();
  15822. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  15823. int64_t cur_level = maybe_layer->layerId();
  15824. if (!isBatchedAtLevel(other, cur_level)) {
  15825. return at::_ops::bitwise_right_shift_Scalar_Tensor::call(self, other);
  15826. }
  15827. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  15828. auto results = batch_rule(self, other_value, other_bdim);
  15829. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  15830. }
  15831. template <typename batch_rule_t, batch_rule_t batch_rule>
  15832. at::Tensor & tril__generated_plumbing(at::Tensor & self, int64_t diagonal) {
  15833. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15834. auto maybe_layer = maybeCurrentDynamicLayer();
  15835. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  15836. int64_t cur_level = maybe_layer->layerId();
  15837. if (!isBatchedAtLevel(self, cur_level)) {
  15838. return at::_ops::tril_::call(self, diagonal);
  15839. }
  15840. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15841. batch_rule(self_value, self_bdim, diagonal);
  15842. return self;
  15843. }
  15844. template <typename batch_rule_t, batch_rule_t batch_rule>
  15845. at::Tensor & triu__generated_plumbing(at::Tensor & self, int64_t diagonal) {
  15846. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15847. auto maybe_layer = maybeCurrentDynamicLayer();
  15848. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  15849. int64_t cur_level = maybe_layer->layerId();
  15850. if (!isBatchedAtLevel(self, cur_level)) {
  15851. return at::_ops::triu_::call(self, diagonal);
  15852. }
  15853. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15854. batch_rule(self_value, self_bdim, diagonal);
  15855. return self;
  15856. }
  15857. template <typename batch_rule_t, batch_rule_t batch_rule>
  15858. at::Tensor & digamma__generated_plumbing(at::Tensor & self) {
  15859. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15860. auto maybe_layer = maybeCurrentDynamicLayer();
  15861. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  15862. int64_t cur_level = maybe_layer->layerId();
  15863. if (!isBatchedAtLevel(self, cur_level)) {
  15864. return at::_ops::digamma_::call(self);
  15865. }
  15866. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15867. batch_rule(self_value, self_bdim);
  15868. return self;
  15869. }
  15870. template <typename batch_rule_t, batch_rule_t batch_rule>
  15871. at::Tensor & lerp__Scalar_generated_plumbing(at::Tensor & self, const at::Tensor & end, const at::Scalar & weight) {
  15872. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15873. auto maybe_layer = maybeCurrentDynamicLayer();
  15874. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  15875. int64_t cur_level = maybe_layer->layerId();
  15876. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(end, cur_level)) {
  15877. return at::_ops::lerp__Scalar::call(self, end, weight);
  15878. }
  15879. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15880. auto [end_value, end_bdim] = unwrapTensorAtLevel(end, cur_level);
  15881. batch_rule(self_value, self_bdim, end_value, end_bdim, weight);
  15882. return self;
  15883. }
  15884. template <typename batch_rule_t, batch_rule_t batch_rule>
  15885. at::Tensor & lerp__Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & end, const at::Tensor & weight) {
  15886. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15887. auto maybe_layer = maybeCurrentDynamicLayer();
  15888. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  15889. int64_t cur_level = maybe_layer->layerId();
  15890. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(end, cur_level) && !isBatchedAtLevel(weight, cur_level)) {
  15891. return at::_ops::lerp__Tensor::call(self, end, weight);
  15892. }
  15893. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15894. auto [end_value, end_bdim] = unwrapTensorAtLevel(end, cur_level);
  15895. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  15896. batch_rule(self_value, self_bdim, end_value, end_bdim, weight_value, weight_bdim);
  15897. return self;
  15898. }
  15899. template <typename batch_rule_t, batch_rule_t batch_rule>
  15900. at::Tensor & addbmm__generated_plumbing(at::Tensor & self, const at::Tensor & batch1, const at::Tensor & batch2, const at::Scalar & beta, const at::Scalar & alpha) {
  15901. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15902. auto maybe_layer = maybeCurrentDynamicLayer();
  15903. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  15904. int64_t cur_level = maybe_layer->layerId();
  15905. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(batch1, cur_level) && !isBatchedAtLevel(batch2, cur_level)) {
  15906. return at::_ops::addbmm_::call(self, batch1, batch2, beta, alpha);
  15907. }
  15908. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15909. auto [batch1_value, batch1_bdim] = unwrapTensorAtLevel(batch1, cur_level);
  15910. auto [batch2_value, batch2_bdim] = unwrapTensorAtLevel(batch2, cur_level);
  15911. batch_rule(self_value, self_bdim, batch1_value, batch1_bdim, batch2_value, batch2_bdim, beta, alpha);
  15912. return self;
  15913. }
  15914. template <typename batch_rule_t, batch_rule_t batch_rule>
  15915. at::Tensor addbmm_generated_plumbing(const at::Tensor & self, const at::Tensor & batch1, const at::Tensor & batch2, const at::Scalar & beta, const at::Scalar & alpha) {
  15916. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15917. auto maybe_layer = maybeCurrentDynamicLayer();
  15918. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  15919. int64_t cur_level = maybe_layer->layerId();
  15920. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(batch1, cur_level) && !isBatchedAtLevel(batch2, cur_level)) {
  15921. return at::_ops::addbmm::call(self, batch1, batch2, beta, alpha);
  15922. }
  15923. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15924. auto [batch1_value, batch1_bdim] = unwrapTensorAtLevel(batch1, cur_level);
  15925. auto [batch2_value, batch2_bdim] = unwrapTensorAtLevel(batch2, cur_level);
  15926. auto results = batch_rule(self_value, self_bdim, batch1_value, batch1_bdim, batch2_value, batch2_bdim, beta, alpha);
  15927. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  15928. }
  15929. template <typename batch_rule_t, batch_rule_t batch_rule>
  15930. at::Tensor & random__from_generated_plumbing(at::Tensor & self, int64_t from, ::std::optional<int64_t> to, ::std::optional<at::Generator> generator) {
  15931. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15932. auto maybe_layer = maybeCurrentDynamicLayer();
  15933. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  15934. int64_t cur_level = maybe_layer->layerId();
  15935. if (!isBatchedAtLevel(self, cur_level)) {
  15936. return at::_ops::random__from::call(self, from, to, generator);
  15937. }
  15938. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15939. batch_rule(self_value, self_bdim, from, to, generator);
  15940. return self;
  15941. }
  15942. template <typename batch_rule_t, batch_rule_t batch_rule>
  15943. at::Tensor & random__to_generated_plumbing(at::Tensor & self, int64_t to, ::std::optional<at::Generator> generator) {
  15944. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15945. auto maybe_layer = maybeCurrentDynamicLayer();
  15946. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  15947. int64_t cur_level = maybe_layer->layerId();
  15948. if (!isBatchedAtLevel(self, cur_level)) {
  15949. return at::_ops::random__to::call(self, to, generator);
  15950. }
  15951. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15952. batch_rule(self_value, self_bdim, to, generator);
  15953. return self;
  15954. }
  15955. template <typename batch_rule_t, batch_rule_t batch_rule>
  15956. at::Tensor & random__generated_plumbing(at::Tensor & self, ::std::optional<at::Generator> generator) {
  15957. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15958. auto maybe_layer = maybeCurrentDynamicLayer();
  15959. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  15960. int64_t cur_level = maybe_layer->layerId();
  15961. if (!isBatchedAtLevel(self, cur_level)) {
  15962. return at::_ops::random_::call(self, generator);
  15963. }
  15964. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15965. batch_rule(self_value, self_bdim, generator);
  15966. return self;
  15967. }
  15968. template <typename batch_rule_t, batch_rule_t batch_rule>
  15969. at::Tensor & uniform__generated_plumbing(at::Tensor & self, double from, double to, ::std::optional<at::Generator> generator) {
  15970. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15971. auto maybe_layer = maybeCurrentDynamicLayer();
  15972. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  15973. int64_t cur_level = maybe_layer->layerId();
  15974. if (!isBatchedAtLevel(self, cur_level)) {
  15975. return at::_ops::uniform_::call(self, from, to, generator);
  15976. }
  15977. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15978. batch_rule(self_value, self_bdim, from, to, generator);
  15979. return self;
  15980. }
  15981. template <typename batch_rule_t, batch_rule_t batch_rule>
  15982. at::Tensor & cauchy__generated_plumbing(at::Tensor & self, double median, double sigma, ::std::optional<at::Generator> generator) {
  15983. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15984. auto maybe_layer = maybeCurrentDynamicLayer();
  15985. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  15986. int64_t cur_level = maybe_layer->layerId();
  15987. if (!isBatchedAtLevel(self, cur_level)) {
  15988. return at::_ops::cauchy_::call(self, median, sigma, generator);
  15989. }
  15990. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  15991. batch_rule(self_value, self_bdim, median, sigma, generator);
  15992. return self;
  15993. }
  15994. template <typename batch_rule_t, batch_rule_t batch_rule>
  15995. at::Tensor & log_normal__generated_plumbing(at::Tensor & self, double mean, double std, ::std::optional<at::Generator> generator) {
  15996. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  15997. auto maybe_layer = maybeCurrentDynamicLayer();
  15998. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  15999. int64_t cur_level = maybe_layer->layerId();
  16000. if (!isBatchedAtLevel(self, cur_level)) {
  16001. return at::_ops::log_normal_::call(self, mean, std, generator);
  16002. }
  16003. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16004. batch_rule(self_value, self_bdim, mean, std, generator);
  16005. return self;
  16006. }
  16007. template <typename batch_rule_t, batch_rule_t batch_rule>
  16008. at::Tensor & exponential__generated_plumbing(at::Tensor & self, double lambd, ::std::optional<at::Generator> generator) {
  16009. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16010. auto maybe_layer = maybeCurrentDynamicLayer();
  16011. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  16012. int64_t cur_level = maybe_layer->layerId();
  16013. if (!isBatchedAtLevel(self, cur_level)) {
  16014. return at::_ops::exponential_::call(self, lambd, generator);
  16015. }
  16016. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16017. batch_rule(self_value, self_bdim, lambd, generator);
  16018. return self;
  16019. }
  16020. template <typename batch_rule_t, batch_rule_t batch_rule>
  16021. at::Tensor & geometric__generated_plumbing(at::Tensor & self, double p, ::std::optional<at::Generator> generator) {
  16022. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16023. auto maybe_layer = maybeCurrentDynamicLayer();
  16024. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  16025. int64_t cur_level = maybe_layer->layerId();
  16026. if (!isBatchedAtLevel(self, cur_level)) {
  16027. return at::_ops::geometric_::call(self, p, generator);
  16028. }
  16029. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16030. batch_rule(self_value, self_bdim, p, generator);
  16031. return self;
  16032. }
  16033. template <typename batch_rule_t, batch_rule_t batch_rule>
  16034. at::Tensor diag_generated_plumbing(const at::Tensor & self, int64_t diagonal) {
  16035. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16036. auto maybe_layer = maybeCurrentDynamicLayer();
  16037. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16038. int64_t cur_level = maybe_layer->layerId();
  16039. if (!isBatchedAtLevel(self, cur_level)) {
  16040. return at::_ops::diag::call(self, diagonal);
  16041. }
  16042. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16043. auto results = batch_rule(self_value, self_bdim, diagonal);
  16044. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16045. }
  16046. template <typename batch_rule_t, batch_rule_t batch_rule>
  16047. at::Tensor cross_generated_plumbing(const at::Tensor & self, const at::Tensor & other, ::std::optional<int64_t> dim) {
  16048. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16049. auto maybe_layer = maybeCurrentDynamicLayer();
  16050. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16051. int64_t cur_level = maybe_layer->layerId();
  16052. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  16053. return at::_ops::cross::call(self, other, dim);
  16054. }
  16055. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16056. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  16057. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim, dim);
  16058. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16059. }
  16060. template <typename batch_rule_t, batch_rule_t batch_rule>
  16061. at::Tensor triu_generated_plumbing(const at::Tensor & self, int64_t diagonal) {
  16062. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16063. auto maybe_layer = maybeCurrentDynamicLayer();
  16064. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16065. int64_t cur_level = maybe_layer->layerId();
  16066. if (!isBatchedAtLevel(self, cur_level)) {
  16067. return at::_ops::triu::call(self, diagonal);
  16068. }
  16069. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16070. auto results = batch_rule(self_value, self_bdim, diagonal);
  16071. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16072. }
  16073. template <typename batch_rule_t, batch_rule_t batch_rule>
  16074. at::Tensor tril_generated_plumbing(const at::Tensor & self, int64_t diagonal) {
  16075. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16076. auto maybe_layer = maybeCurrentDynamicLayer();
  16077. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16078. int64_t cur_level = maybe_layer->layerId();
  16079. if (!isBatchedAtLevel(self, cur_level)) {
  16080. return at::_ops::tril::call(self, diagonal);
  16081. }
  16082. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16083. auto results = batch_rule(self_value, self_bdim, diagonal);
  16084. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16085. }
  16086. template <typename batch_rule_t, batch_rule_t batch_rule>
  16087. at::Tensor trace_generated_plumbing(const at::Tensor & self) {
  16088. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16089. auto maybe_layer = maybeCurrentDynamicLayer();
  16090. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16091. int64_t cur_level = maybe_layer->layerId();
  16092. if (!isBatchedAtLevel(self, cur_level)) {
  16093. return at::_ops::trace::call(self);
  16094. }
  16095. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16096. auto results = batch_rule(self_value, self_bdim);
  16097. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16098. }
  16099. template <typename batch_rule_t, batch_rule_t batch_rule>
  16100. at::Tensor trace_backward_generated_plumbing(const at::Tensor & grad, c10::SymIntArrayRef sizes) {
  16101. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16102. auto maybe_layer = maybeCurrentDynamicLayer();
  16103. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16104. int64_t cur_level = maybe_layer->layerId();
  16105. if (!isBatchedAtLevel(grad, cur_level)) {
  16106. return at::_ops::trace_backward::call(grad, sizes);
  16107. }
  16108. auto [grad_value, grad_bdim] = unwrapTensorAtLevel(grad, cur_level);
  16109. auto results = batch_rule(grad_value, grad_bdim, sizes);
  16110. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16111. }
  16112. template <typename batch_rule_t, batch_rule_t batch_rule>
  16113. at::Tensor ne_Scalar_generated_plumbing(const at::Tensor & self, const at::Scalar & other) {
  16114. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16115. auto maybe_layer = maybeCurrentDynamicLayer();
  16116. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16117. int64_t cur_level = maybe_layer->layerId();
  16118. if (!isBatchedAtLevel(self, cur_level)) {
  16119. return at::_ops::ne_Scalar::call(self, other);
  16120. }
  16121. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16122. auto results = batch_rule(self_value, self_bdim, other);
  16123. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16124. }
  16125. template <typename batch_rule_t, batch_rule_t batch_rule>
  16126. at::Tensor ne_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  16127. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16128. auto maybe_layer = maybeCurrentDynamicLayer();
  16129. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16130. int64_t cur_level = maybe_layer->layerId();
  16131. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  16132. return at::_ops::ne_Tensor::call(self, other);
  16133. }
  16134. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16135. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  16136. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  16137. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16138. }
  16139. template <typename batch_rule_t, batch_rule_t batch_rule>
  16140. at::Tensor & ne__Scalar_generated_plumbing(at::Tensor & self, const at::Scalar & other) {
  16141. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16142. auto maybe_layer = maybeCurrentDynamicLayer();
  16143. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  16144. int64_t cur_level = maybe_layer->layerId();
  16145. if (!isBatchedAtLevel(self, cur_level)) {
  16146. return at::_ops::ne__Scalar::call(self, other);
  16147. }
  16148. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16149. batch_rule(self_value, self_bdim, other);
  16150. return self;
  16151. }
  16152. template <typename batch_rule_t, batch_rule_t batch_rule>
  16153. at::Tensor & ne__Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  16154. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16155. auto maybe_layer = maybeCurrentDynamicLayer();
  16156. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  16157. int64_t cur_level = maybe_layer->layerId();
  16158. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  16159. return at::_ops::ne__Tensor::call(self, other);
  16160. }
  16161. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16162. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  16163. batch_rule(self_value, self_bdim, other_value, other_bdim);
  16164. return self;
  16165. }
  16166. template <typename batch_rule_t, batch_rule_t batch_rule>
  16167. at::Tensor not_equal_Scalar_generated_plumbing(const at::Tensor & self, const at::Scalar & other) {
  16168. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16169. auto maybe_layer = maybeCurrentDynamicLayer();
  16170. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16171. int64_t cur_level = maybe_layer->layerId();
  16172. if (!isBatchedAtLevel(self, cur_level)) {
  16173. return at::_ops::not_equal_Scalar::call(self, other);
  16174. }
  16175. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16176. auto results = batch_rule(self_value, self_bdim, other);
  16177. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16178. }
  16179. template <typename batch_rule_t, batch_rule_t batch_rule>
  16180. at::Tensor not_equal_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  16181. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16182. auto maybe_layer = maybeCurrentDynamicLayer();
  16183. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16184. int64_t cur_level = maybe_layer->layerId();
  16185. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  16186. return at::_ops::not_equal_Tensor::call(self, other);
  16187. }
  16188. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16189. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  16190. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  16191. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16192. }
  16193. template <typename batch_rule_t, batch_rule_t batch_rule>
  16194. at::Tensor & not_equal__Scalar_generated_plumbing(at::Tensor & self, const at::Scalar & other) {
  16195. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16196. auto maybe_layer = maybeCurrentDynamicLayer();
  16197. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  16198. int64_t cur_level = maybe_layer->layerId();
  16199. if (!isBatchedAtLevel(self, cur_level)) {
  16200. return at::_ops::not_equal__Scalar::call(self, other);
  16201. }
  16202. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16203. batch_rule(self_value, self_bdim, other);
  16204. return self;
  16205. }
  16206. template <typename batch_rule_t, batch_rule_t batch_rule>
  16207. at::Tensor & not_equal__Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  16208. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16209. auto maybe_layer = maybeCurrentDynamicLayer();
  16210. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  16211. int64_t cur_level = maybe_layer->layerId();
  16212. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  16213. return at::_ops::not_equal__Tensor::call(self, other);
  16214. }
  16215. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16216. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  16217. batch_rule(self_value, self_bdim, other_value, other_bdim);
  16218. return self;
  16219. }
  16220. template <typename batch_rule_t, batch_rule_t batch_rule>
  16221. at::Tensor eq_Scalar_generated_plumbing(const at::Tensor & self, const at::Scalar & other) {
  16222. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16223. auto maybe_layer = maybeCurrentDynamicLayer();
  16224. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16225. int64_t cur_level = maybe_layer->layerId();
  16226. if (!isBatchedAtLevel(self, cur_level)) {
  16227. return at::_ops::eq_Scalar::call(self, other);
  16228. }
  16229. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16230. auto results = batch_rule(self_value, self_bdim, other);
  16231. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16232. }
  16233. template <typename batch_rule_t, batch_rule_t batch_rule>
  16234. at::Tensor eq_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  16235. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16236. auto maybe_layer = maybeCurrentDynamicLayer();
  16237. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16238. int64_t cur_level = maybe_layer->layerId();
  16239. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  16240. return at::_ops::eq_Tensor::call(self, other);
  16241. }
  16242. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16243. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  16244. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  16245. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16246. }
  16247. template <typename batch_rule_t, batch_rule_t batch_rule>
  16248. at::Tensor ge_Scalar_generated_plumbing(const at::Tensor & self, const at::Scalar & other) {
  16249. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16250. auto maybe_layer = maybeCurrentDynamicLayer();
  16251. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16252. int64_t cur_level = maybe_layer->layerId();
  16253. if (!isBatchedAtLevel(self, cur_level)) {
  16254. return at::_ops::ge_Scalar::call(self, other);
  16255. }
  16256. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16257. auto results = batch_rule(self_value, self_bdim, other);
  16258. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16259. }
  16260. template <typename batch_rule_t, batch_rule_t batch_rule>
  16261. at::Tensor ge_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  16262. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16263. auto maybe_layer = maybeCurrentDynamicLayer();
  16264. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16265. int64_t cur_level = maybe_layer->layerId();
  16266. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  16267. return at::_ops::ge_Tensor::call(self, other);
  16268. }
  16269. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16270. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  16271. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  16272. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16273. }
  16274. template <typename batch_rule_t, batch_rule_t batch_rule>
  16275. at::Tensor & ge__Scalar_generated_plumbing(at::Tensor & self, const at::Scalar & other) {
  16276. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16277. auto maybe_layer = maybeCurrentDynamicLayer();
  16278. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  16279. int64_t cur_level = maybe_layer->layerId();
  16280. if (!isBatchedAtLevel(self, cur_level)) {
  16281. return at::_ops::ge__Scalar::call(self, other);
  16282. }
  16283. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16284. batch_rule(self_value, self_bdim, other);
  16285. return self;
  16286. }
  16287. template <typename batch_rule_t, batch_rule_t batch_rule>
  16288. at::Tensor & ge__Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  16289. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16290. auto maybe_layer = maybeCurrentDynamicLayer();
  16291. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  16292. int64_t cur_level = maybe_layer->layerId();
  16293. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  16294. return at::_ops::ge__Tensor::call(self, other);
  16295. }
  16296. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16297. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  16298. batch_rule(self_value, self_bdim, other_value, other_bdim);
  16299. return self;
  16300. }
  16301. template <typename batch_rule_t, batch_rule_t batch_rule>
  16302. at::Tensor greater_equal_Scalar_generated_plumbing(const at::Tensor & self, const at::Scalar & other) {
  16303. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16304. auto maybe_layer = maybeCurrentDynamicLayer();
  16305. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16306. int64_t cur_level = maybe_layer->layerId();
  16307. if (!isBatchedAtLevel(self, cur_level)) {
  16308. return at::_ops::greater_equal_Scalar::call(self, other);
  16309. }
  16310. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16311. auto results = batch_rule(self_value, self_bdim, other);
  16312. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16313. }
  16314. template <typename batch_rule_t, batch_rule_t batch_rule>
  16315. at::Tensor greater_equal_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  16316. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16317. auto maybe_layer = maybeCurrentDynamicLayer();
  16318. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16319. int64_t cur_level = maybe_layer->layerId();
  16320. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  16321. return at::_ops::greater_equal_Tensor::call(self, other);
  16322. }
  16323. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16324. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  16325. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  16326. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16327. }
  16328. template <typename batch_rule_t, batch_rule_t batch_rule>
  16329. at::Tensor & greater_equal__Scalar_generated_plumbing(at::Tensor & self, const at::Scalar & other) {
  16330. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16331. auto maybe_layer = maybeCurrentDynamicLayer();
  16332. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  16333. int64_t cur_level = maybe_layer->layerId();
  16334. if (!isBatchedAtLevel(self, cur_level)) {
  16335. return at::_ops::greater_equal__Scalar::call(self, other);
  16336. }
  16337. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16338. batch_rule(self_value, self_bdim, other);
  16339. return self;
  16340. }
  16341. template <typename batch_rule_t, batch_rule_t batch_rule>
  16342. at::Tensor & greater_equal__Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  16343. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16344. auto maybe_layer = maybeCurrentDynamicLayer();
  16345. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  16346. int64_t cur_level = maybe_layer->layerId();
  16347. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  16348. return at::_ops::greater_equal__Tensor::call(self, other);
  16349. }
  16350. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16351. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  16352. batch_rule(self_value, self_bdim, other_value, other_bdim);
  16353. return self;
  16354. }
  16355. template <typename batch_rule_t, batch_rule_t batch_rule>
  16356. at::Tensor le_Scalar_generated_plumbing(const at::Tensor & self, const at::Scalar & other) {
  16357. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16358. auto maybe_layer = maybeCurrentDynamicLayer();
  16359. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16360. int64_t cur_level = maybe_layer->layerId();
  16361. if (!isBatchedAtLevel(self, cur_level)) {
  16362. return at::_ops::le_Scalar::call(self, other);
  16363. }
  16364. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16365. auto results = batch_rule(self_value, self_bdim, other);
  16366. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16367. }
  16368. template <typename batch_rule_t, batch_rule_t batch_rule>
  16369. at::Tensor le_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  16370. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16371. auto maybe_layer = maybeCurrentDynamicLayer();
  16372. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16373. int64_t cur_level = maybe_layer->layerId();
  16374. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  16375. return at::_ops::le_Tensor::call(self, other);
  16376. }
  16377. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16378. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  16379. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  16380. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16381. }
  16382. template <typename batch_rule_t, batch_rule_t batch_rule>
  16383. at::Tensor & le__Scalar_generated_plumbing(at::Tensor & self, const at::Scalar & other) {
  16384. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16385. auto maybe_layer = maybeCurrentDynamicLayer();
  16386. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  16387. int64_t cur_level = maybe_layer->layerId();
  16388. if (!isBatchedAtLevel(self, cur_level)) {
  16389. return at::_ops::le__Scalar::call(self, other);
  16390. }
  16391. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16392. batch_rule(self_value, self_bdim, other);
  16393. return self;
  16394. }
  16395. template <typename batch_rule_t, batch_rule_t batch_rule>
  16396. at::Tensor & le__Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  16397. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16398. auto maybe_layer = maybeCurrentDynamicLayer();
  16399. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  16400. int64_t cur_level = maybe_layer->layerId();
  16401. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  16402. return at::_ops::le__Tensor::call(self, other);
  16403. }
  16404. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16405. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  16406. batch_rule(self_value, self_bdim, other_value, other_bdim);
  16407. return self;
  16408. }
  16409. template <typename batch_rule_t, batch_rule_t batch_rule>
  16410. at::Tensor less_equal_Scalar_generated_plumbing(const at::Tensor & self, const at::Scalar & other) {
  16411. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16412. auto maybe_layer = maybeCurrentDynamicLayer();
  16413. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16414. int64_t cur_level = maybe_layer->layerId();
  16415. if (!isBatchedAtLevel(self, cur_level)) {
  16416. return at::_ops::less_equal_Scalar::call(self, other);
  16417. }
  16418. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16419. auto results = batch_rule(self_value, self_bdim, other);
  16420. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16421. }
  16422. template <typename batch_rule_t, batch_rule_t batch_rule>
  16423. at::Tensor less_equal_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  16424. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16425. auto maybe_layer = maybeCurrentDynamicLayer();
  16426. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16427. int64_t cur_level = maybe_layer->layerId();
  16428. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  16429. return at::_ops::less_equal_Tensor::call(self, other);
  16430. }
  16431. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16432. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  16433. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  16434. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16435. }
  16436. template <typename batch_rule_t, batch_rule_t batch_rule>
  16437. at::Tensor & less_equal__Scalar_generated_plumbing(at::Tensor & self, const at::Scalar & other) {
  16438. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16439. auto maybe_layer = maybeCurrentDynamicLayer();
  16440. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  16441. int64_t cur_level = maybe_layer->layerId();
  16442. if (!isBatchedAtLevel(self, cur_level)) {
  16443. return at::_ops::less_equal__Scalar::call(self, other);
  16444. }
  16445. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16446. batch_rule(self_value, self_bdim, other);
  16447. return self;
  16448. }
  16449. template <typename batch_rule_t, batch_rule_t batch_rule>
  16450. at::Tensor & less_equal__Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  16451. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16452. auto maybe_layer = maybeCurrentDynamicLayer();
  16453. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  16454. int64_t cur_level = maybe_layer->layerId();
  16455. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  16456. return at::_ops::less_equal__Tensor::call(self, other);
  16457. }
  16458. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16459. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  16460. batch_rule(self_value, self_bdim, other_value, other_bdim);
  16461. return self;
  16462. }
  16463. template <typename batch_rule_t, batch_rule_t batch_rule>
  16464. at::Tensor gt_Scalar_generated_plumbing(const at::Tensor & self, const at::Scalar & other) {
  16465. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16466. auto maybe_layer = maybeCurrentDynamicLayer();
  16467. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16468. int64_t cur_level = maybe_layer->layerId();
  16469. if (!isBatchedAtLevel(self, cur_level)) {
  16470. return at::_ops::gt_Scalar::call(self, other);
  16471. }
  16472. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16473. auto results = batch_rule(self_value, self_bdim, other);
  16474. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16475. }
  16476. template <typename batch_rule_t, batch_rule_t batch_rule>
  16477. at::Tensor gt_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  16478. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16479. auto maybe_layer = maybeCurrentDynamicLayer();
  16480. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16481. int64_t cur_level = maybe_layer->layerId();
  16482. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  16483. return at::_ops::gt_Tensor::call(self, other);
  16484. }
  16485. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16486. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  16487. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  16488. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16489. }
  16490. template <typename batch_rule_t, batch_rule_t batch_rule>
  16491. at::Tensor & gt__Scalar_generated_plumbing(at::Tensor & self, const at::Scalar & other) {
  16492. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16493. auto maybe_layer = maybeCurrentDynamicLayer();
  16494. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  16495. int64_t cur_level = maybe_layer->layerId();
  16496. if (!isBatchedAtLevel(self, cur_level)) {
  16497. return at::_ops::gt__Scalar::call(self, other);
  16498. }
  16499. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16500. batch_rule(self_value, self_bdim, other);
  16501. return self;
  16502. }
  16503. template <typename batch_rule_t, batch_rule_t batch_rule>
  16504. at::Tensor & gt__Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  16505. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16506. auto maybe_layer = maybeCurrentDynamicLayer();
  16507. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  16508. int64_t cur_level = maybe_layer->layerId();
  16509. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  16510. return at::_ops::gt__Tensor::call(self, other);
  16511. }
  16512. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16513. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  16514. batch_rule(self_value, self_bdim, other_value, other_bdim);
  16515. return self;
  16516. }
  16517. template <typename batch_rule_t, batch_rule_t batch_rule>
  16518. at::Tensor greater_Scalar_generated_plumbing(const at::Tensor & self, const at::Scalar & other) {
  16519. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16520. auto maybe_layer = maybeCurrentDynamicLayer();
  16521. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16522. int64_t cur_level = maybe_layer->layerId();
  16523. if (!isBatchedAtLevel(self, cur_level)) {
  16524. return at::_ops::greater_Scalar::call(self, other);
  16525. }
  16526. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16527. auto results = batch_rule(self_value, self_bdim, other);
  16528. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16529. }
  16530. template <typename batch_rule_t, batch_rule_t batch_rule>
  16531. at::Tensor greater_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  16532. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16533. auto maybe_layer = maybeCurrentDynamicLayer();
  16534. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16535. int64_t cur_level = maybe_layer->layerId();
  16536. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  16537. return at::_ops::greater_Tensor::call(self, other);
  16538. }
  16539. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16540. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  16541. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  16542. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16543. }
  16544. template <typename batch_rule_t, batch_rule_t batch_rule>
  16545. at::Tensor & greater__Scalar_generated_plumbing(at::Tensor & self, const at::Scalar & other) {
  16546. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16547. auto maybe_layer = maybeCurrentDynamicLayer();
  16548. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  16549. int64_t cur_level = maybe_layer->layerId();
  16550. if (!isBatchedAtLevel(self, cur_level)) {
  16551. return at::_ops::greater__Scalar::call(self, other);
  16552. }
  16553. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16554. batch_rule(self_value, self_bdim, other);
  16555. return self;
  16556. }
  16557. template <typename batch_rule_t, batch_rule_t batch_rule>
  16558. at::Tensor & greater__Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  16559. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16560. auto maybe_layer = maybeCurrentDynamicLayer();
  16561. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  16562. int64_t cur_level = maybe_layer->layerId();
  16563. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  16564. return at::_ops::greater__Tensor::call(self, other);
  16565. }
  16566. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16567. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  16568. batch_rule(self_value, self_bdim, other_value, other_bdim);
  16569. return self;
  16570. }
  16571. template <typename batch_rule_t, batch_rule_t batch_rule>
  16572. at::Tensor lt_Scalar_generated_plumbing(const at::Tensor & self, const at::Scalar & other) {
  16573. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16574. auto maybe_layer = maybeCurrentDynamicLayer();
  16575. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16576. int64_t cur_level = maybe_layer->layerId();
  16577. if (!isBatchedAtLevel(self, cur_level)) {
  16578. return at::_ops::lt_Scalar::call(self, other);
  16579. }
  16580. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16581. auto results = batch_rule(self_value, self_bdim, other);
  16582. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16583. }
  16584. template <typename batch_rule_t, batch_rule_t batch_rule>
  16585. at::Tensor lt_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  16586. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16587. auto maybe_layer = maybeCurrentDynamicLayer();
  16588. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16589. int64_t cur_level = maybe_layer->layerId();
  16590. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  16591. return at::_ops::lt_Tensor::call(self, other);
  16592. }
  16593. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16594. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  16595. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  16596. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16597. }
  16598. template <typename batch_rule_t, batch_rule_t batch_rule>
  16599. at::Tensor & lt__Scalar_generated_plumbing(at::Tensor & self, const at::Scalar & other) {
  16600. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16601. auto maybe_layer = maybeCurrentDynamicLayer();
  16602. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  16603. int64_t cur_level = maybe_layer->layerId();
  16604. if (!isBatchedAtLevel(self, cur_level)) {
  16605. return at::_ops::lt__Scalar::call(self, other);
  16606. }
  16607. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16608. batch_rule(self_value, self_bdim, other);
  16609. return self;
  16610. }
  16611. template <typename batch_rule_t, batch_rule_t batch_rule>
  16612. at::Tensor & lt__Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  16613. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16614. auto maybe_layer = maybeCurrentDynamicLayer();
  16615. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  16616. int64_t cur_level = maybe_layer->layerId();
  16617. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  16618. return at::_ops::lt__Tensor::call(self, other);
  16619. }
  16620. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16621. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  16622. batch_rule(self_value, self_bdim, other_value, other_bdim);
  16623. return self;
  16624. }
  16625. template <typename batch_rule_t, batch_rule_t batch_rule>
  16626. at::Tensor less_Scalar_generated_plumbing(const at::Tensor & self, const at::Scalar & other) {
  16627. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16628. auto maybe_layer = maybeCurrentDynamicLayer();
  16629. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16630. int64_t cur_level = maybe_layer->layerId();
  16631. if (!isBatchedAtLevel(self, cur_level)) {
  16632. return at::_ops::less_Scalar::call(self, other);
  16633. }
  16634. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16635. auto results = batch_rule(self_value, self_bdim, other);
  16636. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16637. }
  16638. template <typename batch_rule_t, batch_rule_t batch_rule>
  16639. at::Tensor less_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  16640. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16641. auto maybe_layer = maybeCurrentDynamicLayer();
  16642. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16643. int64_t cur_level = maybe_layer->layerId();
  16644. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  16645. return at::_ops::less_Tensor::call(self, other);
  16646. }
  16647. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16648. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  16649. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  16650. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16651. }
  16652. template <typename batch_rule_t, batch_rule_t batch_rule>
  16653. at::Tensor & less__Scalar_generated_plumbing(at::Tensor & self, const at::Scalar & other) {
  16654. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16655. auto maybe_layer = maybeCurrentDynamicLayer();
  16656. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  16657. int64_t cur_level = maybe_layer->layerId();
  16658. if (!isBatchedAtLevel(self, cur_level)) {
  16659. return at::_ops::less__Scalar::call(self, other);
  16660. }
  16661. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16662. batch_rule(self_value, self_bdim, other);
  16663. return self;
  16664. }
  16665. template <typename batch_rule_t, batch_rule_t batch_rule>
  16666. at::Tensor & less__Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  16667. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16668. auto maybe_layer = maybeCurrentDynamicLayer();
  16669. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  16670. int64_t cur_level = maybe_layer->layerId();
  16671. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  16672. return at::_ops::less__Tensor::call(self, other);
  16673. }
  16674. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16675. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  16676. batch_rule(self_value, self_bdim, other_value, other_bdim);
  16677. return self;
  16678. }
  16679. template <typename batch_rule_t, batch_rule_t batch_rule>
  16680. at::Tensor take_generated_plumbing(const at::Tensor & self, const at::Tensor & index) {
  16681. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16682. auto maybe_layer = maybeCurrentDynamicLayer();
  16683. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16684. int64_t cur_level = maybe_layer->layerId();
  16685. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(index, cur_level)) {
  16686. return at::_ops::take::call(self, index);
  16687. }
  16688. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16689. auto [index_value, index_bdim] = unwrapTensorAtLevel(index, cur_level);
  16690. auto results = batch_rule(self_value, self_bdim, index_value, index_bdim);
  16691. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16692. }
  16693. template <typename batch_rule_t, batch_rule_t batch_rule>
  16694. at::Tensor take_along_dim_generated_plumbing(const at::Tensor & self, const at::Tensor & indices, ::std::optional<int64_t> dim) {
  16695. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16696. auto maybe_layer = maybeCurrentDynamicLayer();
  16697. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16698. int64_t cur_level = maybe_layer->layerId();
  16699. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(indices, cur_level)) {
  16700. return at::_ops::take_along_dim::call(self, indices, dim);
  16701. }
  16702. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16703. auto [indices_value, indices_bdim] = unwrapTensorAtLevel(indices, cur_level);
  16704. auto results = batch_rule(self_value, self_bdim, indices_value, indices_bdim, dim);
  16705. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16706. }
  16707. template <typename batch_rule_t, batch_rule_t batch_rule>
  16708. at::Tensor index_select_generated_plumbing(const at::Tensor & self, int64_t dim, const at::Tensor & index) {
  16709. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16710. auto maybe_layer = maybeCurrentDynamicLayer();
  16711. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16712. int64_t cur_level = maybe_layer->layerId();
  16713. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(index, cur_level)) {
  16714. return at::_ops::index_select::call(self, dim, index);
  16715. }
  16716. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16717. auto [index_value, index_bdim] = unwrapTensorAtLevel(index, cur_level);
  16718. auto results = batch_rule(self_value, self_bdim, dim, index_value, index_bdim);
  16719. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16720. }
  16721. template <typename batch_rule_t, batch_rule_t batch_rule>
  16722. at::Tensor index_select_dimname_generated_plumbing(const at::Tensor & self, at::Dimname dim, const at::Tensor & index) {
  16723. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16724. auto maybe_layer = maybeCurrentDynamicLayer();
  16725. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16726. int64_t cur_level = maybe_layer->layerId();
  16727. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(index, cur_level)) {
  16728. return at::_ops::index_select_dimname::call(self, dim, index);
  16729. }
  16730. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16731. auto [index_value, index_bdim] = unwrapTensorAtLevel(index, cur_level);
  16732. auto results = batch_rule(self_value, self_bdim, dim, index_value, index_bdim);
  16733. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16734. }
  16735. template <typename batch_rule_t, batch_rule_t batch_rule>
  16736. at::Tensor index_select_backward_generated_plumbing(const at::Tensor & grad, c10::SymIntArrayRef self_sizes, int64_t dim, const at::Tensor & index) {
  16737. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16738. auto maybe_layer = maybeCurrentDynamicLayer();
  16739. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16740. int64_t cur_level = maybe_layer->layerId();
  16741. if (!isBatchedAtLevel(grad, cur_level) && !isBatchedAtLevel(index, cur_level)) {
  16742. return at::_ops::index_select_backward::call(grad, self_sizes, dim, index);
  16743. }
  16744. auto [grad_value, grad_bdim] = unwrapTensorAtLevel(grad, cur_level);
  16745. auto [index_value, index_bdim] = unwrapTensorAtLevel(index, cur_level);
  16746. auto results = batch_rule(grad_value, grad_bdim, self_sizes, dim, index_value, index_bdim);
  16747. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16748. }
  16749. template <typename batch_rule_t, batch_rule_t batch_rule>
  16750. at::Tensor masked_select_generated_plumbing(const at::Tensor & self, const at::Tensor & mask) {
  16751. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16752. auto maybe_layer = maybeCurrentDynamicLayer();
  16753. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16754. int64_t cur_level = maybe_layer->layerId();
  16755. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(mask, cur_level)) {
  16756. return at::_ops::masked_select::call(self, mask);
  16757. }
  16758. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16759. auto [mask_value, mask_bdim] = unwrapTensorAtLevel(mask, cur_level);
  16760. auto results = batch_rule(self_value, self_bdim, mask_value, mask_bdim);
  16761. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16762. }
  16763. template <typename batch_rule_t, batch_rule_t batch_rule>
  16764. at::Tensor masked_select_backward_generated_plumbing(const at::Tensor & grad, const at::Tensor & input, const at::Tensor & mask) {
  16765. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16766. auto maybe_layer = maybeCurrentDynamicLayer();
  16767. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16768. int64_t cur_level = maybe_layer->layerId();
  16769. if (!isBatchedAtLevel(grad, cur_level) && !isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(mask, cur_level)) {
  16770. return at::_ops::masked_select_backward::call(grad, input, mask);
  16771. }
  16772. auto [grad_value, grad_bdim] = unwrapTensorAtLevel(grad, cur_level);
  16773. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  16774. auto [mask_value, mask_bdim] = unwrapTensorAtLevel(mask, cur_level);
  16775. auto results = batch_rule(grad_value, grad_bdim, input_value, input_bdim, mask_value, mask_bdim);
  16776. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16777. }
  16778. template <typename batch_rule_t, batch_rule_t batch_rule>
  16779. at::Tensor nonzero_generated_plumbing(const at::Tensor & self) {
  16780. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16781. auto maybe_layer = maybeCurrentDynamicLayer();
  16782. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16783. int64_t cur_level = maybe_layer->layerId();
  16784. if (!isBatchedAtLevel(self, cur_level)) {
  16785. return at::_ops::nonzero::call(self);
  16786. }
  16787. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16788. auto results = batch_rule(self_value, self_bdim);
  16789. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16790. }
  16791. template <typename batch_rule_t, batch_rule_t batch_rule>
  16792. at::Tensor nonzero_static_generated_plumbing(const at::Tensor & self, c10::SymInt size, int64_t fill_value) {
  16793. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16794. auto maybe_layer = maybeCurrentDynamicLayer();
  16795. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16796. int64_t cur_level = maybe_layer->layerId();
  16797. if (!isBatchedAtLevel(self, cur_level)) {
  16798. return at::_ops::nonzero_static::call(self, size, fill_value);
  16799. }
  16800. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16801. auto results = batch_rule(self_value, self_bdim, size, fill_value);
  16802. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16803. }
  16804. template <typename batch_rule_t, batch_rule_t batch_rule>
  16805. ::std::vector<at::Tensor> nonzero_numpy_generated_plumbing(const at::Tensor & self) {
  16806. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16807. auto maybe_layer = maybeCurrentDynamicLayer();
  16808. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16809. int64_t cur_level = maybe_layer->layerId();
  16810. if (!isBatchedAtLevel(self, cur_level)) {
  16811. return at::_ops::nonzero_numpy::call(self);
  16812. }
  16813. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16814. auto results = batch_rule(self_value, self_bdim);
  16815. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  16816. }
  16817. template <typename batch_rule_t, batch_rule_t batch_rule>
  16818. at::Tensor argwhere_generated_plumbing(const at::Tensor & self) {
  16819. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16820. auto maybe_layer = maybeCurrentDynamicLayer();
  16821. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16822. int64_t cur_level = maybe_layer->layerId();
  16823. if (!isBatchedAtLevel(self, cur_level)) {
  16824. return at::_ops::argwhere::call(self);
  16825. }
  16826. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16827. auto results = batch_rule(self_value, self_bdim);
  16828. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16829. }
  16830. template <typename batch_rule_t, batch_rule_t batch_rule>
  16831. at::Tensor gather_generated_plumbing(const at::Tensor & self, int64_t dim, const at::Tensor & index, bool sparse_grad) {
  16832. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16833. auto maybe_layer = maybeCurrentDynamicLayer();
  16834. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16835. int64_t cur_level = maybe_layer->layerId();
  16836. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(index, cur_level)) {
  16837. return at::_ops::gather::call(self, dim, index, sparse_grad);
  16838. }
  16839. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16840. auto [index_value, index_bdim] = unwrapTensorAtLevel(index, cur_level);
  16841. auto results = batch_rule(self_value, self_bdim, dim, index_value, index_bdim, sparse_grad);
  16842. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16843. }
  16844. template <typename batch_rule_t, batch_rule_t batch_rule>
  16845. at::Tensor gather_backward_generated_plumbing(const at::Tensor & grad, const at::Tensor & self, int64_t dim, const at::Tensor & index, bool sparse_grad) {
  16846. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16847. auto maybe_layer = maybeCurrentDynamicLayer();
  16848. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16849. int64_t cur_level = maybe_layer->layerId();
  16850. if (!isBatchedAtLevel(grad, cur_level) && !isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(index, cur_level)) {
  16851. return at::_ops::gather_backward::call(grad, self, dim, index, sparse_grad);
  16852. }
  16853. auto [grad_value, grad_bdim] = unwrapTensorAtLevel(grad, cur_level);
  16854. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16855. auto [index_value, index_bdim] = unwrapTensorAtLevel(index, cur_level);
  16856. auto results = batch_rule(grad_value, grad_bdim, self_value, self_bdim, dim, index_value, index_bdim, sparse_grad);
  16857. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16858. }
  16859. template <typename batch_rule_t, batch_rule_t batch_rule>
  16860. at::Tensor gather_dimname_generated_plumbing(const at::Tensor & self, at::Dimname dim, const at::Tensor & index, bool sparse_grad) {
  16861. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16862. auto maybe_layer = maybeCurrentDynamicLayer();
  16863. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16864. int64_t cur_level = maybe_layer->layerId();
  16865. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(index, cur_level)) {
  16866. return at::_ops::gather_dimname::call(self, dim, index, sparse_grad);
  16867. }
  16868. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16869. auto [index_value, index_bdim] = unwrapTensorAtLevel(index, cur_level);
  16870. auto results = batch_rule(self_value, self_bdim, dim, index_value, index_bdim, sparse_grad);
  16871. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16872. }
  16873. template <typename batch_rule_t, batch_rule_t batch_rule>
  16874. at::Tensor _gather_sparse_backward_generated_plumbing(const at::Tensor & self, int64_t dim, const at::Tensor & index, const at::Tensor & grad) {
  16875. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16876. auto maybe_layer = maybeCurrentDynamicLayer();
  16877. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16878. int64_t cur_level = maybe_layer->layerId();
  16879. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(index, cur_level) && !isBatchedAtLevel(grad, cur_level)) {
  16880. return at::_ops::_gather_sparse_backward::call(self, dim, index, grad);
  16881. }
  16882. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16883. auto [index_value, index_bdim] = unwrapTensorAtLevel(index, cur_level);
  16884. auto [grad_value, grad_bdim] = unwrapTensorAtLevel(grad, cur_level);
  16885. auto results = batch_rule(self_value, self_bdim, dim, index_value, index_bdim, grad_value, grad_bdim);
  16886. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16887. }
  16888. template <typename batch_rule_t, batch_rule_t batch_rule>
  16889. at::Tensor addcmul_generated_plumbing(const at::Tensor & self, const at::Tensor & tensor1, const at::Tensor & tensor2, const at::Scalar & value) {
  16890. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16891. auto maybe_layer = maybeCurrentDynamicLayer();
  16892. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16893. int64_t cur_level = maybe_layer->layerId();
  16894. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(tensor1, cur_level) && !isBatchedAtLevel(tensor2, cur_level)) {
  16895. return at::_ops::addcmul::call(self, tensor1, tensor2, value);
  16896. }
  16897. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16898. auto [tensor1_value, tensor1_bdim] = unwrapTensorAtLevel(tensor1, cur_level);
  16899. auto [tensor2_value, tensor2_bdim] = unwrapTensorAtLevel(tensor2, cur_level);
  16900. auto results = batch_rule(self_value, self_bdim, tensor1_value, tensor1_bdim, tensor2_value, tensor2_bdim, value);
  16901. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16902. }
  16903. template <typename batch_rule_t, batch_rule_t batch_rule>
  16904. at::Tensor & addcmul__generated_plumbing(at::Tensor & self, const at::Tensor & tensor1, const at::Tensor & tensor2, const at::Scalar & value) {
  16905. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16906. auto maybe_layer = maybeCurrentDynamicLayer();
  16907. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  16908. int64_t cur_level = maybe_layer->layerId();
  16909. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(tensor1, cur_level) && !isBatchedAtLevel(tensor2, cur_level)) {
  16910. return at::_ops::addcmul_::call(self, tensor1, tensor2, value);
  16911. }
  16912. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16913. auto [tensor1_value, tensor1_bdim] = unwrapTensorAtLevel(tensor1, cur_level);
  16914. auto [tensor2_value, tensor2_bdim] = unwrapTensorAtLevel(tensor2, cur_level);
  16915. batch_rule(self_value, self_bdim, tensor1_value, tensor1_bdim, tensor2_value, tensor2_bdim, value);
  16916. return self;
  16917. }
  16918. template <typename batch_rule_t, batch_rule_t batch_rule>
  16919. at::Tensor addcdiv_generated_plumbing(const at::Tensor & self, const at::Tensor & tensor1, const at::Tensor & tensor2, const at::Scalar & value) {
  16920. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16921. auto maybe_layer = maybeCurrentDynamicLayer();
  16922. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16923. int64_t cur_level = maybe_layer->layerId();
  16924. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(tensor1, cur_level) && !isBatchedAtLevel(tensor2, cur_level)) {
  16925. return at::_ops::addcdiv::call(self, tensor1, tensor2, value);
  16926. }
  16927. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16928. auto [tensor1_value, tensor1_bdim] = unwrapTensorAtLevel(tensor1, cur_level);
  16929. auto [tensor2_value, tensor2_bdim] = unwrapTensorAtLevel(tensor2, cur_level);
  16930. auto results = batch_rule(self_value, self_bdim, tensor1_value, tensor1_bdim, tensor2_value, tensor2_bdim, value);
  16931. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16932. }
  16933. template <typename batch_rule_t, batch_rule_t batch_rule>
  16934. at::Tensor & addcdiv__generated_plumbing(at::Tensor & self, const at::Tensor & tensor1, const at::Tensor & tensor2, const at::Scalar & value) {
  16935. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16936. auto maybe_layer = maybeCurrentDynamicLayer();
  16937. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  16938. int64_t cur_level = maybe_layer->layerId();
  16939. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(tensor1, cur_level) && !isBatchedAtLevel(tensor2, cur_level)) {
  16940. return at::_ops::addcdiv_::call(self, tensor1, tensor2, value);
  16941. }
  16942. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16943. auto [tensor1_value, tensor1_bdim] = unwrapTensorAtLevel(tensor1, cur_level);
  16944. auto [tensor2_value, tensor2_bdim] = unwrapTensorAtLevel(tensor2, cur_level);
  16945. batch_rule(self_value, self_bdim, tensor1_value, tensor1_bdim, tensor2_value, tensor2_bdim, value);
  16946. return self;
  16947. }
  16948. template <typename batch_rule_t, batch_rule_t batch_rule>
  16949. at::Tensor cross_entropy_loss_generated_plumbing(const at::Tensor & self, const at::Tensor & target, const ::std::optional<at::Tensor> & weight, int64_t reduction, c10::SymInt ignore_index, double label_smoothing) {
  16950. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16951. auto maybe_layer = maybeCurrentDynamicLayer();
  16952. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16953. int64_t cur_level = maybe_layer->layerId();
  16954. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(target, cur_level) && !isBatchedAtLevel(weight, cur_level)) {
  16955. return at::_ops::cross_entropy_loss::call(self, target, weight, reduction, ignore_index, label_smoothing);
  16956. }
  16957. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16958. auto [target_value, target_bdim] = unwrapTensorAtLevel(target, cur_level);
  16959. std::optional<Tensor> weight_value;
  16960. std::optional<int64_t> weight_bdim;
  16961. if (weight) {
  16962. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  16963. }
  16964. auto results = batch_rule(self_value, self_bdim, target_value, target_bdim, weight_value, weight_bdim, reduction, ignore_index, label_smoothing);
  16965. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  16966. }
  16967. template <typename batch_rule_t, batch_rule_t batch_rule>
  16968. ::std::tuple<at::Tensor,at::Tensor> triangular_solve_generated_plumbing(const at::Tensor & self, const at::Tensor & A, bool upper, bool transpose, bool unitriangular) {
  16969. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16970. auto maybe_layer = maybeCurrentDynamicLayer();
  16971. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16972. int64_t cur_level = maybe_layer->layerId();
  16973. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(A, cur_level)) {
  16974. return at::_ops::triangular_solve::call(self, A, upper, transpose, unitriangular);
  16975. }
  16976. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  16977. auto [A_value, A_bdim] = unwrapTensorAtLevel(A, cur_level);
  16978. auto results = batch_rule(self_value, self_bdim, A_value, A_bdim, upper, transpose, unitriangular);
  16979. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  16980. }
  16981. template <typename batch_rule_t, batch_rule_t batch_rule>
  16982. void _linalg_check_errors_generated_plumbing(const at::Tensor & info, c10::string_view api_name, bool is_matrix) {
  16983. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16984. auto maybe_layer = maybeCurrentDynamicLayer();
  16985. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  16986. int64_t cur_level = maybe_layer->layerId();
  16987. if (!isBatchedAtLevel(info, cur_level)) {
  16988. return at::_ops::_linalg_check_errors::call(info, api_name, is_matrix);
  16989. }
  16990. auto [info_value, info_bdim] = unwrapTensorAtLevel(info, cur_level);
  16991. batch_rule(info_value, info_bdim, api_name, is_matrix);
  16992. }
  16993. template <typename batch_rule_t, batch_rule_t batch_rule>
  16994. at::Tensor linalg_solve_triangular_generated_plumbing(const at::Tensor & self, const at::Tensor & B, bool upper, bool left, bool unitriangular) {
  16995. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  16996. auto maybe_layer = maybeCurrentDynamicLayer();
  16997. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  16998. int64_t cur_level = maybe_layer->layerId();
  16999. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(B, cur_level)) {
  17000. return at::_ops::linalg_solve_triangular::call(self, B, upper, left, unitriangular);
  17001. }
  17002. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17003. auto [B_value, B_bdim] = unwrapTensorAtLevel(B, cur_level);
  17004. auto results = batch_rule(self_value, self_bdim, B_value, B_bdim, upper, left, unitriangular);
  17005. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17006. }
  17007. template <typename batch_rule_t, batch_rule_t batch_rule>
  17008. at::Tensor linalg_vander_generated_plumbing(const at::Tensor & x, ::std::optional<c10::SymInt> N) {
  17009. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17010. auto maybe_layer = maybeCurrentDynamicLayer();
  17011. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17012. int64_t cur_level = maybe_layer->layerId();
  17013. if (!isBatchedAtLevel(x, cur_level)) {
  17014. return at::_ops::linalg_vander::call(x, N);
  17015. }
  17016. auto [x_value, x_bdim] = unwrapTensorAtLevel(x, cur_level);
  17017. auto results = batch_rule(x_value, x_bdim, N);
  17018. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17019. }
  17020. template <typename batch_rule_t, batch_rule_t batch_rule>
  17021. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> svd_generated_plumbing(const at::Tensor & self, bool some, bool compute_uv) {
  17022. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17023. auto maybe_layer = maybeCurrentDynamicLayer();
  17024. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17025. int64_t cur_level = maybe_layer->layerId();
  17026. if (!isBatchedAtLevel(self, cur_level)) {
  17027. return at::_ops::svd::call(self, some, compute_uv);
  17028. }
  17029. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17030. auto results = batch_rule(self_value, self_bdim, some, compute_uv);
  17031. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  17032. }
  17033. template <typename batch_rule_t, batch_rule_t batch_rule>
  17034. at::Tensor swapaxes_generated_plumbing(const at::Tensor & self, int64_t axis0, int64_t axis1) {
  17035. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17036. auto maybe_layer = maybeCurrentDynamicLayer();
  17037. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17038. int64_t cur_level = maybe_layer->layerId();
  17039. if (!isBatchedAtLevel(self, cur_level)) {
  17040. return at::_ops::swapaxes::call(self, axis0, axis1);
  17041. }
  17042. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17043. auto results = batch_rule(self_value, self_bdim, axis0, axis1);
  17044. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17045. }
  17046. template <typename batch_rule_t, batch_rule_t batch_rule>
  17047. at::Tensor swapdims_generated_plumbing(const at::Tensor & self, int64_t dim0, int64_t dim1) {
  17048. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17049. auto maybe_layer = maybeCurrentDynamicLayer();
  17050. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17051. int64_t cur_level = maybe_layer->layerId();
  17052. if (!isBatchedAtLevel(self, cur_level)) {
  17053. return at::_ops::swapdims::call(self, dim0, dim1);
  17054. }
  17055. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17056. auto results = batch_rule(self_value, self_bdim, dim0, dim1);
  17057. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17058. }
  17059. template <typename batch_rule_t, batch_rule_t batch_rule>
  17060. at::Tensor cholesky_generated_plumbing(const at::Tensor & self, bool upper) {
  17061. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17062. auto maybe_layer = maybeCurrentDynamicLayer();
  17063. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17064. int64_t cur_level = maybe_layer->layerId();
  17065. if (!isBatchedAtLevel(self, cur_level)) {
  17066. return at::_ops::cholesky::call(self, upper);
  17067. }
  17068. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17069. auto results = batch_rule(self_value, self_bdim, upper);
  17070. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17071. }
  17072. template <typename batch_rule_t, batch_rule_t batch_rule>
  17073. at::Tensor cholesky_solve_generated_plumbing(const at::Tensor & self, const at::Tensor & input2, bool upper) {
  17074. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17075. auto maybe_layer = maybeCurrentDynamicLayer();
  17076. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17077. int64_t cur_level = maybe_layer->layerId();
  17078. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(input2, cur_level)) {
  17079. return at::_ops::cholesky_solve::call(self, input2, upper);
  17080. }
  17081. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17082. auto [input2_value, input2_bdim] = unwrapTensorAtLevel(input2, cur_level);
  17083. auto results = batch_rule(self_value, self_bdim, input2_value, input2_bdim, upper);
  17084. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17085. }
  17086. template <typename batch_rule_t, batch_rule_t batch_rule>
  17087. at::Tensor _cholesky_solve_helper_generated_plumbing(const at::Tensor & self, const at::Tensor & A, bool upper) {
  17088. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17089. auto maybe_layer = maybeCurrentDynamicLayer();
  17090. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17091. int64_t cur_level = maybe_layer->layerId();
  17092. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(A, cur_level)) {
  17093. return at::_ops::_cholesky_solve_helper::call(self, A, upper);
  17094. }
  17095. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17096. auto [A_value, A_bdim] = unwrapTensorAtLevel(A, cur_level);
  17097. auto results = batch_rule(self_value, self_bdim, A_value, A_bdim, upper);
  17098. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17099. }
  17100. template <typename batch_rule_t, batch_rule_t batch_rule>
  17101. at::Tensor cholesky_inverse_generated_plumbing(const at::Tensor & self, bool upper) {
  17102. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17103. auto maybe_layer = maybeCurrentDynamicLayer();
  17104. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17105. int64_t cur_level = maybe_layer->layerId();
  17106. if (!isBatchedAtLevel(self, cur_level)) {
  17107. return at::_ops::cholesky_inverse::call(self, upper);
  17108. }
  17109. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17110. auto results = batch_rule(self_value, self_bdim, upper);
  17111. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17112. }
  17113. template <typename batch_rule_t, batch_rule_t batch_rule>
  17114. ::std::tuple<at::Tensor,at::Tensor> qr_generated_plumbing(const at::Tensor & self, bool some) {
  17115. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17116. auto maybe_layer = maybeCurrentDynamicLayer();
  17117. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17118. int64_t cur_level = maybe_layer->layerId();
  17119. if (!isBatchedAtLevel(self, cur_level)) {
  17120. return at::_ops::qr::call(self, some);
  17121. }
  17122. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17123. auto results = batch_rule(self_value, self_bdim, some);
  17124. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  17125. }
  17126. template <typename batch_rule_t, batch_rule_t batch_rule>
  17127. ::std::tuple<at::Tensor,at::Tensor> geqrf_generated_plumbing(const at::Tensor & self) {
  17128. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17129. auto maybe_layer = maybeCurrentDynamicLayer();
  17130. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17131. int64_t cur_level = maybe_layer->layerId();
  17132. if (!isBatchedAtLevel(self, cur_level)) {
  17133. return at::_ops::geqrf::call(self);
  17134. }
  17135. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17136. auto results = batch_rule(self_value, self_bdim);
  17137. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  17138. }
  17139. template <typename batch_rule_t, batch_rule_t batch_rule>
  17140. at::Tensor orgqr_generated_plumbing(const at::Tensor & self, const at::Tensor & input2) {
  17141. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17142. auto maybe_layer = maybeCurrentDynamicLayer();
  17143. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17144. int64_t cur_level = maybe_layer->layerId();
  17145. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(input2, cur_level)) {
  17146. return at::_ops::orgqr::call(self, input2);
  17147. }
  17148. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17149. auto [input2_value, input2_bdim] = unwrapTensorAtLevel(input2, cur_level);
  17150. auto results = batch_rule(self_value, self_bdim, input2_value, input2_bdim);
  17151. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17152. }
  17153. template <typename batch_rule_t, batch_rule_t batch_rule>
  17154. at::Tensor ormqr_generated_plumbing(const at::Tensor & self, const at::Tensor & input2, const at::Tensor & input3, bool left, bool transpose) {
  17155. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17156. auto maybe_layer = maybeCurrentDynamicLayer();
  17157. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17158. int64_t cur_level = maybe_layer->layerId();
  17159. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(input2, cur_level) && !isBatchedAtLevel(input3, cur_level)) {
  17160. return at::_ops::ormqr::call(self, input2, input3, left, transpose);
  17161. }
  17162. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17163. auto [input2_value, input2_bdim] = unwrapTensorAtLevel(input2, cur_level);
  17164. auto [input3_value, input3_bdim] = unwrapTensorAtLevel(input3, cur_level);
  17165. auto results = batch_rule(self_value, self_bdim, input2_value, input2_bdim, input3_value, input3_bdim, left, transpose);
  17166. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17167. }
  17168. template <typename batch_rule_t, batch_rule_t batch_rule>
  17169. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> _lu_with_info_generated_plumbing(const at::Tensor & self, bool pivot, bool check_errors) {
  17170. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17171. auto maybe_layer = maybeCurrentDynamicLayer();
  17172. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17173. int64_t cur_level = maybe_layer->layerId();
  17174. if (!isBatchedAtLevel(self, cur_level)) {
  17175. return at::_ops::_lu_with_info::call(self, pivot, check_errors);
  17176. }
  17177. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17178. auto results = batch_rule(self_value, self_bdim, pivot, check_errors);
  17179. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  17180. }
  17181. template <typename batch_rule_t, batch_rule_t batch_rule>
  17182. at::Tensor lu_solve_generated_plumbing(const at::Tensor & self, const at::Tensor & LU_data, const at::Tensor & LU_pivots) {
  17183. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17184. auto maybe_layer = maybeCurrentDynamicLayer();
  17185. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17186. int64_t cur_level = maybe_layer->layerId();
  17187. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(LU_data, cur_level) && !isBatchedAtLevel(LU_pivots, cur_level)) {
  17188. return at::_ops::lu_solve::call(self, LU_data, LU_pivots);
  17189. }
  17190. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17191. auto [LU_data_value, LU_data_bdim] = unwrapTensorAtLevel(LU_data, cur_level);
  17192. auto [LU_pivots_value, LU_pivots_bdim] = unwrapTensorAtLevel(LU_pivots, cur_level);
  17193. auto results = batch_rule(self_value, self_bdim, LU_data_value, LU_data_bdim, LU_pivots_value, LU_pivots_bdim);
  17194. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17195. }
  17196. template <typename batch_rule_t, batch_rule_t batch_rule>
  17197. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> lu_unpack_generated_plumbing(const at::Tensor & LU_data, const at::Tensor & LU_pivots, bool unpack_data, bool unpack_pivots) {
  17198. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17199. auto maybe_layer = maybeCurrentDynamicLayer();
  17200. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17201. int64_t cur_level = maybe_layer->layerId();
  17202. if (!isBatchedAtLevel(LU_data, cur_level) && !isBatchedAtLevel(LU_pivots, cur_level)) {
  17203. return at::_ops::lu_unpack::call(LU_data, LU_pivots, unpack_data, unpack_pivots);
  17204. }
  17205. auto [LU_data_value, LU_data_bdim] = unwrapTensorAtLevel(LU_data, cur_level);
  17206. auto [LU_pivots_value, LU_pivots_bdim] = unwrapTensorAtLevel(LU_pivots, cur_level);
  17207. auto results = batch_rule(LU_data_value, LU_data_bdim, LU_pivots_value, LU_pivots_bdim, unpack_data, unpack_pivots);
  17208. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  17209. }
  17210. template <typename batch_rule_t, batch_rule_t batch_rule>
  17211. at::Tensor multinomial_generated_plumbing(const at::Tensor & self, c10::SymInt num_samples, bool replacement, ::std::optional<at::Generator> generator) {
  17212. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17213. auto maybe_layer = maybeCurrentDynamicLayer();
  17214. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17215. int64_t cur_level = maybe_layer->layerId();
  17216. if (!isBatchedAtLevel(self, cur_level)) {
  17217. return at::_ops::multinomial::call(self, num_samples, replacement, generator);
  17218. }
  17219. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17220. auto results = batch_rule(self_value, self_bdim, num_samples, replacement, generator);
  17221. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17222. }
  17223. template <typename batch_rule_t, batch_rule_t batch_rule>
  17224. at::Tensor & lgamma__generated_plumbing(at::Tensor & self) {
  17225. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17226. auto maybe_layer = maybeCurrentDynamicLayer();
  17227. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  17228. int64_t cur_level = maybe_layer->layerId();
  17229. if (!isBatchedAtLevel(self, cur_level)) {
  17230. return at::_ops::lgamma_::call(self);
  17231. }
  17232. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17233. batch_rule(self_value, self_bdim);
  17234. return self;
  17235. }
  17236. template <typename batch_rule_t, batch_rule_t batch_rule>
  17237. at::Tensor lgamma_generated_plumbing(const at::Tensor & self) {
  17238. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17239. auto maybe_layer = maybeCurrentDynamicLayer();
  17240. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17241. int64_t cur_level = maybe_layer->layerId();
  17242. if (!isBatchedAtLevel(self, cur_level)) {
  17243. return at::_ops::lgamma::call(self);
  17244. }
  17245. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17246. auto results = batch_rule(self_value, self_bdim);
  17247. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17248. }
  17249. template <typename batch_rule_t, batch_rule_t batch_rule>
  17250. at::Tensor digamma_generated_plumbing(const at::Tensor & self) {
  17251. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17252. auto maybe_layer = maybeCurrentDynamicLayer();
  17253. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17254. int64_t cur_level = maybe_layer->layerId();
  17255. if (!isBatchedAtLevel(self, cur_level)) {
  17256. return at::_ops::digamma::call(self);
  17257. }
  17258. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17259. auto results = batch_rule(self_value, self_bdim);
  17260. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17261. }
  17262. template <typename batch_rule_t, batch_rule_t batch_rule>
  17263. at::Tensor polygamma_generated_plumbing(int64_t n, const at::Tensor & self) {
  17264. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17265. auto maybe_layer = maybeCurrentDynamicLayer();
  17266. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17267. int64_t cur_level = maybe_layer->layerId();
  17268. if (!isBatchedAtLevel(self, cur_level)) {
  17269. return at::_ops::polygamma::call(n, self);
  17270. }
  17271. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17272. auto results = batch_rule(n, self_value, self_bdim);
  17273. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17274. }
  17275. template <typename batch_rule_t, batch_rule_t batch_rule>
  17276. at::Tensor & polygamma__generated_plumbing(at::Tensor & self, int64_t n) {
  17277. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17278. auto maybe_layer = maybeCurrentDynamicLayer();
  17279. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  17280. int64_t cur_level = maybe_layer->layerId();
  17281. if (!isBatchedAtLevel(self, cur_level)) {
  17282. return at::_ops::polygamma_::call(self, n);
  17283. }
  17284. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17285. batch_rule(self_value, self_bdim, n);
  17286. return self;
  17287. }
  17288. template <typename batch_rule_t, batch_rule_t batch_rule>
  17289. at::Tensor erfinv_generated_plumbing(const at::Tensor & self) {
  17290. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17291. auto maybe_layer = maybeCurrentDynamicLayer();
  17292. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17293. int64_t cur_level = maybe_layer->layerId();
  17294. if (!isBatchedAtLevel(self, cur_level)) {
  17295. return at::_ops::erfinv::call(self);
  17296. }
  17297. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17298. auto results = batch_rule(self_value, self_bdim);
  17299. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17300. }
  17301. template <typename batch_rule_t, batch_rule_t batch_rule>
  17302. at::Tensor & erfinv__generated_plumbing(at::Tensor & self) {
  17303. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17304. auto maybe_layer = maybeCurrentDynamicLayer();
  17305. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  17306. int64_t cur_level = maybe_layer->layerId();
  17307. if (!isBatchedAtLevel(self, cur_level)) {
  17308. return at::_ops::erfinv_::call(self);
  17309. }
  17310. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17311. batch_rule(self_value, self_bdim);
  17312. return self;
  17313. }
  17314. template <typename batch_rule_t, batch_rule_t batch_rule>
  17315. at::Tensor i0_generated_plumbing(const at::Tensor & self) {
  17316. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17317. auto maybe_layer = maybeCurrentDynamicLayer();
  17318. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17319. int64_t cur_level = maybe_layer->layerId();
  17320. if (!isBatchedAtLevel(self, cur_level)) {
  17321. return at::_ops::i0::call(self);
  17322. }
  17323. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17324. auto results = batch_rule(self_value, self_bdim);
  17325. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17326. }
  17327. template <typename batch_rule_t, batch_rule_t batch_rule>
  17328. at::Tensor & i0__generated_plumbing(at::Tensor & self) {
  17329. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17330. auto maybe_layer = maybeCurrentDynamicLayer();
  17331. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  17332. int64_t cur_level = maybe_layer->layerId();
  17333. if (!isBatchedAtLevel(self, cur_level)) {
  17334. return at::_ops::i0_::call(self);
  17335. }
  17336. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17337. batch_rule(self_value, self_bdim);
  17338. return self;
  17339. }
  17340. template <typename batch_rule_t, batch_rule_t batch_rule>
  17341. at::Tensor sign_generated_plumbing(const at::Tensor & self) {
  17342. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17343. auto maybe_layer = maybeCurrentDynamicLayer();
  17344. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17345. int64_t cur_level = maybe_layer->layerId();
  17346. if (!isBatchedAtLevel(self, cur_level)) {
  17347. return at::_ops::sign::call(self);
  17348. }
  17349. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17350. auto results = batch_rule(self_value, self_bdim);
  17351. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17352. }
  17353. template <typename batch_rule_t, batch_rule_t batch_rule>
  17354. at::Tensor & sign__generated_plumbing(at::Tensor & self) {
  17355. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17356. auto maybe_layer = maybeCurrentDynamicLayer();
  17357. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  17358. int64_t cur_level = maybe_layer->layerId();
  17359. if (!isBatchedAtLevel(self, cur_level)) {
  17360. return at::_ops::sign_::call(self);
  17361. }
  17362. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17363. batch_rule(self_value, self_bdim);
  17364. return self;
  17365. }
  17366. template <typename batch_rule_t, batch_rule_t batch_rule>
  17367. at::Tensor signbit_generated_plumbing(const at::Tensor & self) {
  17368. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17369. auto maybe_layer = maybeCurrentDynamicLayer();
  17370. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17371. int64_t cur_level = maybe_layer->layerId();
  17372. if (!isBatchedAtLevel(self, cur_level)) {
  17373. return at::_ops::signbit::call(self);
  17374. }
  17375. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17376. auto results = batch_rule(self_value, self_bdim);
  17377. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17378. }
  17379. template <typename batch_rule_t, batch_rule_t batch_rule>
  17380. at::Tensor dist_generated_plumbing(const at::Tensor & self, const at::Tensor & other, const at::Scalar & p) {
  17381. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17382. auto maybe_layer = maybeCurrentDynamicLayer();
  17383. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17384. int64_t cur_level = maybe_layer->layerId();
  17385. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  17386. return at::_ops::dist::call(self, other, p);
  17387. }
  17388. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17389. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  17390. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim, p);
  17391. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17392. }
  17393. template <typename batch_rule_t, batch_rule_t batch_rule>
  17394. at::Tensor & atan2__generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  17395. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17396. auto maybe_layer = maybeCurrentDynamicLayer();
  17397. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  17398. int64_t cur_level = maybe_layer->layerId();
  17399. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  17400. return at::_ops::atan2_::call(self, other);
  17401. }
  17402. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17403. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  17404. batch_rule(self_value, self_bdim, other_value, other_bdim);
  17405. return self;
  17406. }
  17407. template <typename batch_rule_t, batch_rule_t batch_rule>
  17408. at::Tensor atan2_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  17409. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17410. auto maybe_layer = maybeCurrentDynamicLayer();
  17411. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17412. int64_t cur_level = maybe_layer->layerId();
  17413. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  17414. return at::_ops::atan2::call(self, other);
  17415. }
  17416. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17417. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  17418. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  17419. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17420. }
  17421. template <typename batch_rule_t, batch_rule_t batch_rule>
  17422. at::Tensor arctan2_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  17423. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17424. auto maybe_layer = maybeCurrentDynamicLayer();
  17425. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17426. int64_t cur_level = maybe_layer->layerId();
  17427. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  17428. return at::_ops::arctan2::call(self, other);
  17429. }
  17430. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17431. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  17432. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  17433. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17434. }
  17435. template <typename batch_rule_t, batch_rule_t batch_rule>
  17436. at::Tensor & arctan2__generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  17437. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17438. auto maybe_layer = maybeCurrentDynamicLayer();
  17439. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  17440. int64_t cur_level = maybe_layer->layerId();
  17441. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  17442. return at::_ops::arctan2_::call(self, other);
  17443. }
  17444. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17445. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  17446. batch_rule(self_value, self_bdim, other_value, other_bdim);
  17447. return self;
  17448. }
  17449. template <typename batch_rule_t, batch_rule_t batch_rule>
  17450. at::Tensor lerp_Scalar_generated_plumbing(const at::Tensor & self, const at::Tensor & end, const at::Scalar & weight) {
  17451. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17452. auto maybe_layer = maybeCurrentDynamicLayer();
  17453. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17454. int64_t cur_level = maybe_layer->layerId();
  17455. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(end, cur_level)) {
  17456. return at::_ops::lerp_Scalar::call(self, end, weight);
  17457. }
  17458. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17459. auto [end_value, end_bdim] = unwrapTensorAtLevel(end, cur_level);
  17460. auto results = batch_rule(self_value, self_bdim, end_value, end_bdim, weight);
  17461. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17462. }
  17463. template <typename batch_rule_t, batch_rule_t batch_rule>
  17464. at::Tensor lerp_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & end, const at::Tensor & weight) {
  17465. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17466. auto maybe_layer = maybeCurrentDynamicLayer();
  17467. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17468. int64_t cur_level = maybe_layer->layerId();
  17469. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(end, cur_level) && !isBatchedAtLevel(weight, cur_level)) {
  17470. return at::_ops::lerp_Tensor::call(self, end, weight);
  17471. }
  17472. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17473. auto [end_value, end_bdim] = unwrapTensorAtLevel(end, cur_level);
  17474. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  17475. auto results = batch_rule(self_value, self_bdim, end_value, end_bdim, weight_value, weight_bdim);
  17476. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17477. }
  17478. template <typename batch_rule_t, batch_rule_t batch_rule>
  17479. at::Tensor histc_generated_plumbing(const at::Tensor & self, int64_t bins, const at::Scalar & min, const at::Scalar & max) {
  17480. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17481. auto maybe_layer = maybeCurrentDynamicLayer();
  17482. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17483. int64_t cur_level = maybe_layer->layerId();
  17484. if (!isBatchedAtLevel(self, cur_level)) {
  17485. return at::_ops::histc::call(self, bins, min, max);
  17486. }
  17487. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17488. auto results = batch_rule(self_value, self_bdim, bins, min, max);
  17489. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17490. }
  17491. template <typename batch_rule_t, batch_rule_t batch_rule>
  17492. ::std::tuple<at::Tensor,at::Tensor> histogram_bins_tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & bins, const ::std::optional<at::Tensor> & weight, bool density) {
  17493. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17494. auto maybe_layer = maybeCurrentDynamicLayer();
  17495. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17496. int64_t cur_level = maybe_layer->layerId();
  17497. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(bins, cur_level) && !isBatchedAtLevel(weight, cur_level)) {
  17498. return at::_ops::histogram_bins_tensor::call(self, bins, weight, density);
  17499. }
  17500. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17501. auto [bins_value, bins_bdim] = unwrapTensorAtLevel(bins, cur_level);
  17502. std::optional<Tensor> weight_value;
  17503. std::optional<int64_t> weight_bdim;
  17504. if (weight) {
  17505. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  17506. }
  17507. auto results = batch_rule(self_value, self_bdim, bins_value, bins_bdim, weight_value, weight_bdim, density);
  17508. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  17509. }
  17510. template <typename batch_rule_t, batch_rule_t batch_rule>
  17511. ::std::tuple<at::Tensor,at::Tensor> histogram_bin_ct_generated_plumbing(const at::Tensor & self, int64_t bins, ::std::optional<at::ArrayRef<double>> range, const ::std::optional<at::Tensor> & weight, bool density) {
  17512. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17513. auto maybe_layer = maybeCurrentDynamicLayer();
  17514. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17515. int64_t cur_level = maybe_layer->layerId();
  17516. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(weight, cur_level)) {
  17517. return at::_ops::histogram_bin_ct::call(self, bins, range, weight, density);
  17518. }
  17519. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17520. std::optional<Tensor> weight_value;
  17521. std::optional<int64_t> weight_bdim;
  17522. if (weight) {
  17523. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  17524. }
  17525. auto results = batch_rule(self_value, self_bdim, bins, range, weight_value, weight_bdim, density);
  17526. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  17527. }
  17528. template <typename batch_rule_t, batch_rule_t batch_rule>
  17529. ::std::vector<at::Tensor> _histogramdd_bin_edges_generated_plumbing(const at::Tensor & self, at::IntArrayRef bins, ::std::optional<at::ArrayRef<double>> range, const ::std::optional<at::Tensor> & weight, bool density) {
  17530. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17531. auto maybe_layer = maybeCurrentDynamicLayer();
  17532. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17533. int64_t cur_level = maybe_layer->layerId();
  17534. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(weight, cur_level)) {
  17535. return at::_ops::_histogramdd_bin_edges::call(self, bins, range, weight, density);
  17536. }
  17537. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17538. std::optional<Tensor> weight_value;
  17539. std::optional<int64_t> weight_bdim;
  17540. if (weight) {
  17541. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  17542. }
  17543. auto results = batch_rule(self_value, self_bdim, bins, range, weight_value, weight_bdim, density);
  17544. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  17545. }
  17546. template <typename batch_rule_t, batch_rule_t batch_rule>
  17547. at::Tensor _histogramdd_from_bin_cts_generated_plumbing(const at::Tensor & self, at::IntArrayRef bins, ::std::optional<at::ArrayRef<double>> range, const ::std::optional<at::Tensor> & weight, bool density) {
  17548. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17549. auto maybe_layer = maybeCurrentDynamicLayer();
  17550. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17551. int64_t cur_level = maybe_layer->layerId();
  17552. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(weight, cur_level)) {
  17553. return at::_ops::_histogramdd_from_bin_cts::call(self, bins, range, weight, density);
  17554. }
  17555. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17556. std::optional<Tensor> weight_value;
  17557. std::optional<int64_t> weight_bdim;
  17558. if (weight) {
  17559. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  17560. }
  17561. auto results = batch_rule(self_value, self_bdim, bins, range, weight_value, weight_bdim, density);
  17562. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17563. }
  17564. template <typename batch_rule_t, batch_rule_t batch_rule>
  17565. at::Tensor _histogramdd_from_bin_tensors_generated_plumbing(const at::Tensor & self, at::TensorList bins, const ::std::optional<at::Tensor> & weight, bool density) {
  17566. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17567. auto maybe_layer = maybeCurrentDynamicLayer();
  17568. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17569. int64_t cur_level = maybe_layer->layerId();
  17570. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(bins, cur_level) && !isBatchedAtLevel(weight, cur_level)) {
  17571. return at::_ops::_histogramdd_from_bin_tensors::call(self, bins, weight, density);
  17572. }
  17573. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17574. std::optional<Tensor> weight_value;
  17575. std::optional<int64_t> weight_bdim;
  17576. if (weight) {
  17577. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  17578. }
  17579. auto results = batch_rule(self_value, self_bdim, bins, weight_value, weight_bdim, density);
  17580. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17581. }
  17582. template <typename batch_rule_t, batch_rule_t batch_rule>
  17583. ::std::tuple<at::Tensor,::std::vector<at::Tensor>> histogramdd_generated_plumbing(const at::Tensor & self, at::IntArrayRef bins, ::std::optional<at::ArrayRef<double>> range, const ::std::optional<at::Tensor> & weight, bool density) {
  17584. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17585. auto maybe_layer = maybeCurrentDynamicLayer();
  17586. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17587. int64_t cur_level = maybe_layer->layerId();
  17588. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(weight, cur_level)) {
  17589. return at::_ops::histogramdd::call(self, bins, range, weight, density);
  17590. }
  17591. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17592. std::optional<Tensor> weight_value;
  17593. std::optional<int64_t> weight_bdim;
  17594. if (weight) {
  17595. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  17596. }
  17597. auto results = batch_rule(self_value, self_bdim, bins, range, weight_value, weight_bdim, density);
  17598. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatchedVector(std::get<2>(results), std::get<3>(results), cur_level));
  17599. }
  17600. template <typename batch_rule_t, batch_rule_t batch_rule>
  17601. ::std::tuple<at::Tensor,::std::vector<at::Tensor>> histogramdd_int_bins_generated_plumbing(const at::Tensor & self, int64_t bins, ::std::optional<at::ArrayRef<double>> range, const ::std::optional<at::Tensor> & weight, bool density) {
  17602. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17603. auto maybe_layer = maybeCurrentDynamicLayer();
  17604. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17605. int64_t cur_level = maybe_layer->layerId();
  17606. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(weight, cur_level)) {
  17607. return at::_ops::histogramdd_int_bins::call(self, bins, range, weight, density);
  17608. }
  17609. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17610. std::optional<Tensor> weight_value;
  17611. std::optional<int64_t> weight_bdim;
  17612. if (weight) {
  17613. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  17614. }
  17615. auto results = batch_rule(self_value, self_bdim, bins, range, weight_value, weight_bdim, density);
  17616. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatchedVector(std::get<2>(results), std::get<3>(results), cur_level));
  17617. }
  17618. template <typename batch_rule_t, batch_rule_t batch_rule>
  17619. ::std::tuple<at::Tensor,::std::vector<at::Tensor>> histogramdd_TensorList_bins_generated_plumbing(const at::Tensor & self, at::TensorList bins, ::std::optional<at::ArrayRef<double>> range, const ::std::optional<at::Tensor> & weight, bool density) {
  17620. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17621. auto maybe_layer = maybeCurrentDynamicLayer();
  17622. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17623. int64_t cur_level = maybe_layer->layerId();
  17624. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(bins, cur_level) && !isBatchedAtLevel(weight, cur_level)) {
  17625. return at::_ops::histogramdd_TensorList_bins::call(self, bins, range, weight, density);
  17626. }
  17627. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17628. std::optional<Tensor> weight_value;
  17629. std::optional<int64_t> weight_bdim;
  17630. if (weight) {
  17631. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  17632. }
  17633. auto results = batch_rule(self_value, self_bdim, bins, range, weight_value, weight_bdim, density);
  17634. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatchedVector(std::get<2>(results), std::get<3>(results), cur_level));
  17635. }
  17636. template <typename batch_rule_t, batch_rule_t batch_rule>
  17637. at::Tensor fmod_Scalar_generated_plumbing(const at::Tensor & self, const at::Scalar & other) {
  17638. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17639. auto maybe_layer = maybeCurrentDynamicLayer();
  17640. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17641. int64_t cur_level = maybe_layer->layerId();
  17642. if (!isBatchedAtLevel(self, cur_level)) {
  17643. return at::_ops::fmod_Scalar::call(self, other);
  17644. }
  17645. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17646. auto results = batch_rule(self_value, self_bdim, other);
  17647. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17648. }
  17649. template <typename batch_rule_t, batch_rule_t batch_rule>
  17650. at::Tensor & fmod__Scalar_generated_plumbing(at::Tensor & self, const at::Scalar & other) {
  17651. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17652. auto maybe_layer = maybeCurrentDynamicLayer();
  17653. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  17654. int64_t cur_level = maybe_layer->layerId();
  17655. if (!isBatchedAtLevel(self, cur_level)) {
  17656. return at::_ops::fmod__Scalar::call(self, other);
  17657. }
  17658. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17659. batch_rule(self_value, self_bdim, other);
  17660. return self;
  17661. }
  17662. template <typename batch_rule_t, batch_rule_t batch_rule>
  17663. at::Tensor fmod_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  17664. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17665. auto maybe_layer = maybeCurrentDynamicLayer();
  17666. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17667. int64_t cur_level = maybe_layer->layerId();
  17668. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  17669. return at::_ops::fmod_Tensor::call(self, other);
  17670. }
  17671. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17672. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  17673. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  17674. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17675. }
  17676. template <typename batch_rule_t, batch_rule_t batch_rule>
  17677. at::Tensor & fmod__Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  17678. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17679. auto maybe_layer = maybeCurrentDynamicLayer();
  17680. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  17681. int64_t cur_level = maybe_layer->layerId();
  17682. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  17683. return at::_ops::fmod__Tensor::call(self, other);
  17684. }
  17685. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17686. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  17687. batch_rule(self_value, self_bdim, other_value, other_bdim);
  17688. return self;
  17689. }
  17690. template <typename batch_rule_t, batch_rule_t batch_rule>
  17691. at::Tensor hypot_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  17692. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17693. auto maybe_layer = maybeCurrentDynamicLayer();
  17694. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17695. int64_t cur_level = maybe_layer->layerId();
  17696. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  17697. return at::_ops::hypot::call(self, other);
  17698. }
  17699. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17700. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  17701. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  17702. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17703. }
  17704. template <typename batch_rule_t, batch_rule_t batch_rule>
  17705. at::Tensor & hypot__generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  17706. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17707. auto maybe_layer = maybeCurrentDynamicLayer();
  17708. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  17709. int64_t cur_level = maybe_layer->layerId();
  17710. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  17711. return at::_ops::hypot_::call(self, other);
  17712. }
  17713. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17714. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  17715. batch_rule(self_value, self_bdim, other_value, other_bdim);
  17716. return self;
  17717. }
  17718. template <typename batch_rule_t, batch_rule_t batch_rule>
  17719. at::Tensor igamma_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  17720. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17721. auto maybe_layer = maybeCurrentDynamicLayer();
  17722. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17723. int64_t cur_level = maybe_layer->layerId();
  17724. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  17725. return at::_ops::igamma::call(self, other);
  17726. }
  17727. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17728. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  17729. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  17730. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17731. }
  17732. template <typename batch_rule_t, batch_rule_t batch_rule>
  17733. at::Tensor & igamma__generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  17734. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17735. auto maybe_layer = maybeCurrentDynamicLayer();
  17736. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  17737. int64_t cur_level = maybe_layer->layerId();
  17738. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  17739. return at::_ops::igamma_::call(self, other);
  17740. }
  17741. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17742. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  17743. batch_rule(self_value, self_bdim, other_value, other_bdim);
  17744. return self;
  17745. }
  17746. template <typename batch_rule_t, batch_rule_t batch_rule>
  17747. at::Tensor igammac_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  17748. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17749. auto maybe_layer = maybeCurrentDynamicLayer();
  17750. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17751. int64_t cur_level = maybe_layer->layerId();
  17752. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  17753. return at::_ops::igammac::call(self, other);
  17754. }
  17755. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17756. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  17757. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  17758. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17759. }
  17760. template <typename batch_rule_t, batch_rule_t batch_rule>
  17761. at::Tensor & igammac__generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  17762. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17763. auto maybe_layer = maybeCurrentDynamicLayer();
  17764. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  17765. int64_t cur_level = maybe_layer->layerId();
  17766. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  17767. return at::_ops::igammac_::call(self, other);
  17768. }
  17769. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17770. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  17771. batch_rule(self_value, self_bdim, other_value, other_bdim);
  17772. return self;
  17773. }
  17774. template <typename batch_rule_t, batch_rule_t batch_rule>
  17775. at::Tensor nextafter_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  17776. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17777. auto maybe_layer = maybeCurrentDynamicLayer();
  17778. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17779. int64_t cur_level = maybe_layer->layerId();
  17780. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  17781. return at::_ops::nextafter::call(self, other);
  17782. }
  17783. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17784. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  17785. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  17786. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17787. }
  17788. template <typename batch_rule_t, batch_rule_t batch_rule>
  17789. at::Tensor & nextafter__generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  17790. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17791. auto maybe_layer = maybeCurrentDynamicLayer();
  17792. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  17793. int64_t cur_level = maybe_layer->layerId();
  17794. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  17795. return at::_ops::nextafter_::call(self, other);
  17796. }
  17797. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17798. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  17799. batch_rule(self_value, self_bdim, other_value, other_bdim);
  17800. return self;
  17801. }
  17802. template <typename batch_rule_t, batch_rule_t batch_rule>
  17803. at::Tensor remainder_Scalar_generated_plumbing(const at::Tensor & self, const at::Scalar & other) {
  17804. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17805. auto maybe_layer = maybeCurrentDynamicLayer();
  17806. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17807. int64_t cur_level = maybe_layer->layerId();
  17808. if (!isBatchedAtLevel(self, cur_level)) {
  17809. return at::_ops::remainder_Scalar::call(self, other);
  17810. }
  17811. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17812. auto results = batch_rule(self_value, self_bdim, other);
  17813. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17814. }
  17815. template <typename batch_rule_t, batch_rule_t batch_rule>
  17816. at::Tensor & remainder__Scalar_generated_plumbing(at::Tensor & self, const at::Scalar & other) {
  17817. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17818. auto maybe_layer = maybeCurrentDynamicLayer();
  17819. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  17820. int64_t cur_level = maybe_layer->layerId();
  17821. if (!isBatchedAtLevel(self, cur_level)) {
  17822. return at::_ops::remainder__Scalar::call(self, other);
  17823. }
  17824. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17825. batch_rule(self_value, self_bdim, other);
  17826. return self;
  17827. }
  17828. template <typename batch_rule_t, batch_rule_t batch_rule>
  17829. at::Tensor remainder_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  17830. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17831. auto maybe_layer = maybeCurrentDynamicLayer();
  17832. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17833. int64_t cur_level = maybe_layer->layerId();
  17834. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  17835. return at::_ops::remainder_Tensor::call(self, other);
  17836. }
  17837. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17838. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  17839. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  17840. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17841. }
  17842. template <typename batch_rule_t, batch_rule_t batch_rule>
  17843. at::Tensor & remainder__Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & other) {
  17844. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17845. auto maybe_layer = maybeCurrentDynamicLayer();
  17846. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  17847. int64_t cur_level = maybe_layer->layerId();
  17848. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  17849. return at::_ops::remainder__Tensor::call(self, other);
  17850. }
  17851. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17852. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  17853. batch_rule(self_value, self_bdim, other_value, other_bdim);
  17854. return self;
  17855. }
  17856. template <typename batch_rule_t, batch_rule_t batch_rule>
  17857. at::Tensor remainder_Scalar_Tensor_generated_plumbing(const at::Scalar & self, const at::Tensor & other) {
  17858. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17859. auto maybe_layer = maybeCurrentDynamicLayer();
  17860. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17861. int64_t cur_level = maybe_layer->layerId();
  17862. if (!isBatchedAtLevel(other, cur_level)) {
  17863. return at::_ops::remainder_Scalar_Tensor::call(self, other);
  17864. }
  17865. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  17866. auto results = batch_rule(self, other_value, other_bdim);
  17867. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17868. }
  17869. template <typename batch_rule_t, batch_rule_t batch_rule>
  17870. at::Tensor min_generated_plumbing(const at::Tensor & self) {
  17871. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17872. auto maybe_layer = maybeCurrentDynamicLayer();
  17873. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17874. int64_t cur_level = maybe_layer->layerId();
  17875. if (!isBatchedAtLevel(self, cur_level)) {
  17876. return at::_ops::min::call(self);
  17877. }
  17878. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17879. auto results = batch_rule(self_value, self_bdim);
  17880. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17881. }
  17882. template <typename batch_rule_t, batch_rule_t batch_rule>
  17883. at::Tensor fmin_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  17884. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17885. auto maybe_layer = maybeCurrentDynamicLayer();
  17886. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17887. int64_t cur_level = maybe_layer->layerId();
  17888. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  17889. return at::_ops::fmin::call(self, other);
  17890. }
  17891. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17892. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  17893. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  17894. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17895. }
  17896. template <typename batch_rule_t, batch_rule_t batch_rule>
  17897. at::Tensor max_generated_plumbing(const at::Tensor & self) {
  17898. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17899. auto maybe_layer = maybeCurrentDynamicLayer();
  17900. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17901. int64_t cur_level = maybe_layer->layerId();
  17902. if (!isBatchedAtLevel(self, cur_level)) {
  17903. return at::_ops::max::call(self);
  17904. }
  17905. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17906. auto results = batch_rule(self_value, self_bdim);
  17907. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17908. }
  17909. template <typename batch_rule_t, batch_rule_t batch_rule>
  17910. at::Tensor fmax_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  17911. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17912. auto maybe_layer = maybeCurrentDynamicLayer();
  17913. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17914. int64_t cur_level = maybe_layer->layerId();
  17915. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  17916. return at::_ops::fmax::call(self, other);
  17917. }
  17918. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17919. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  17920. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  17921. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17922. }
  17923. template <typename batch_rule_t, batch_rule_t batch_rule>
  17924. at::Tensor maximum_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  17925. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17926. auto maybe_layer = maybeCurrentDynamicLayer();
  17927. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17928. int64_t cur_level = maybe_layer->layerId();
  17929. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  17930. return at::_ops::maximum::call(self, other);
  17931. }
  17932. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17933. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  17934. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  17935. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17936. }
  17937. template <typename batch_rule_t, batch_rule_t batch_rule>
  17938. at::Tensor max_other_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  17939. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17940. auto maybe_layer = maybeCurrentDynamicLayer();
  17941. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17942. int64_t cur_level = maybe_layer->layerId();
  17943. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  17944. return at::_ops::max_other::call(self, other);
  17945. }
  17946. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17947. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  17948. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  17949. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17950. }
  17951. template <typename batch_rule_t, batch_rule_t batch_rule>
  17952. at::Tensor minimum_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  17953. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17954. auto maybe_layer = maybeCurrentDynamicLayer();
  17955. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17956. int64_t cur_level = maybe_layer->layerId();
  17957. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  17958. return at::_ops::minimum::call(self, other);
  17959. }
  17960. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17961. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  17962. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  17963. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17964. }
  17965. template <typename batch_rule_t, batch_rule_t batch_rule>
  17966. at::Tensor min_other_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  17967. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17968. auto maybe_layer = maybeCurrentDynamicLayer();
  17969. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17970. int64_t cur_level = maybe_layer->layerId();
  17971. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  17972. return at::_ops::min_other::call(self, other);
  17973. }
  17974. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17975. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  17976. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  17977. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17978. }
  17979. template <typename batch_rule_t, batch_rule_t batch_rule>
  17980. at::Tensor quantile_generated_plumbing(const at::Tensor & self, const at::Tensor & q, ::std::optional<int64_t> dim, bool keepdim, c10::string_view interpolation) {
  17981. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17982. auto maybe_layer = maybeCurrentDynamicLayer();
  17983. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17984. int64_t cur_level = maybe_layer->layerId();
  17985. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(q, cur_level)) {
  17986. return at::_ops::quantile::call(self, q, dim, keepdim, interpolation);
  17987. }
  17988. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  17989. auto [q_value, q_bdim] = unwrapTensorAtLevel(q, cur_level);
  17990. auto results = batch_rule(self_value, self_bdim, q_value, q_bdim, dim, keepdim, interpolation);
  17991. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  17992. }
  17993. template <typename batch_rule_t, batch_rule_t batch_rule>
  17994. at::Tensor quantile_scalar_generated_plumbing(const at::Tensor & self, double q, ::std::optional<int64_t> dim, bool keepdim, c10::string_view interpolation) {
  17995. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  17996. auto maybe_layer = maybeCurrentDynamicLayer();
  17997. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  17998. int64_t cur_level = maybe_layer->layerId();
  17999. if (!isBatchedAtLevel(self, cur_level)) {
  18000. return at::_ops::quantile_scalar::call(self, q, dim, keepdim, interpolation);
  18001. }
  18002. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  18003. auto results = batch_rule(self_value, self_bdim, q, dim, keepdim, interpolation);
  18004. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  18005. }
  18006. template <typename batch_rule_t, batch_rule_t batch_rule>
  18007. at::Tensor nanquantile_generated_plumbing(const at::Tensor & self, const at::Tensor & q, ::std::optional<int64_t> dim, bool keepdim, c10::string_view interpolation) {
  18008. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18009. auto maybe_layer = maybeCurrentDynamicLayer();
  18010. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18011. int64_t cur_level = maybe_layer->layerId();
  18012. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(q, cur_level)) {
  18013. return at::_ops::nanquantile::call(self, q, dim, keepdim, interpolation);
  18014. }
  18015. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  18016. auto [q_value, q_bdim] = unwrapTensorAtLevel(q, cur_level);
  18017. auto results = batch_rule(self_value, self_bdim, q_value, q_bdim, dim, keepdim, interpolation);
  18018. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  18019. }
  18020. template <typename batch_rule_t, batch_rule_t batch_rule>
  18021. at::Tensor nanquantile_scalar_generated_plumbing(const at::Tensor & self, double q, ::std::optional<int64_t> dim, bool keepdim, c10::string_view interpolation) {
  18022. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18023. auto maybe_layer = maybeCurrentDynamicLayer();
  18024. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18025. int64_t cur_level = maybe_layer->layerId();
  18026. if (!isBatchedAtLevel(self, cur_level)) {
  18027. return at::_ops::nanquantile_scalar::call(self, q, dim, keepdim, interpolation);
  18028. }
  18029. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  18030. auto results = batch_rule(self_value, self_bdim, q, dim, keepdim, interpolation);
  18031. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  18032. }
  18033. template <typename batch_rule_t, batch_rule_t batch_rule>
  18034. ::std::tuple<at::Tensor,at::Tensor> sort_generated_plumbing(const at::Tensor & self, int64_t dim, bool descending) {
  18035. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18036. auto maybe_layer = maybeCurrentDynamicLayer();
  18037. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18038. int64_t cur_level = maybe_layer->layerId();
  18039. if (!isBatchedAtLevel(self, cur_level)) {
  18040. return at::_ops::sort::call(self, dim, descending);
  18041. }
  18042. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  18043. auto results = batch_rule(self_value, self_bdim, dim, descending);
  18044. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  18045. }
  18046. template <typename batch_rule_t, batch_rule_t batch_rule>
  18047. ::std::tuple<at::Tensor,at::Tensor> sort_stable_generated_plumbing(const at::Tensor & self, ::std::optional<bool> stable, int64_t dim, bool descending) {
  18048. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18049. auto maybe_layer = maybeCurrentDynamicLayer();
  18050. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18051. int64_t cur_level = maybe_layer->layerId();
  18052. if (!isBatchedAtLevel(self, cur_level)) {
  18053. return at::_ops::sort_stable::call(self, stable, dim, descending);
  18054. }
  18055. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  18056. auto results = batch_rule(self_value, self_bdim, stable, dim, descending);
  18057. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  18058. }
  18059. template <typename batch_rule_t, batch_rule_t batch_rule>
  18060. ::std::tuple<at::Tensor,at::Tensor> sort_dimname_generated_plumbing(const at::Tensor & self, at::Dimname dim, bool descending) {
  18061. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18062. auto maybe_layer = maybeCurrentDynamicLayer();
  18063. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18064. int64_t cur_level = maybe_layer->layerId();
  18065. if (!isBatchedAtLevel(self, cur_level)) {
  18066. return at::_ops::sort_dimname::call(self, dim, descending);
  18067. }
  18068. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  18069. auto results = batch_rule(self_value, self_bdim, dim, descending);
  18070. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  18071. }
  18072. template <typename batch_rule_t, batch_rule_t batch_rule>
  18073. ::std::tuple<at::Tensor,at::Tensor> sort_dimname_stable_generated_plumbing(const at::Tensor & self, ::std::optional<bool> stable, at::Dimname dim, bool descending) {
  18074. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18075. auto maybe_layer = maybeCurrentDynamicLayer();
  18076. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18077. int64_t cur_level = maybe_layer->layerId();
  18078. if (!isBatchedAtLevel(self, cur_level)) {
  18079. return at::_ops::sort_dimname_stable::call(self, stable, dim, descending);
  18080. }
  18081. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  18082. auto results = batch_rule(self_value, self_bdim, stable, dim, descending);
  18083. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  18084. }
  18085. template <typename batch_rule_t, batch_rule_t batch_rule>
  18086. at::Tensor msort_generated_plumbing(const at::Tensor & self) {
  18087. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18088. auto maybe_layer = maybeCurrentDynamicLayer();
  18089. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18090. int64_t cur_level = maybe_layer->layerId();
  18091. if (!isBatchedAtLevel(self, cur_level)) {
  18092. return at::_ops::msort::call(self);
  18093. }
  18094. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  18095. auto results = batch_rule(self_value, self_bdim);
  18096. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  18097. }
  18098. template <typename batch_rule_t, batch_rule_t batch_rule>
  18099. at::Tensor argsort_generated_plumbing(const at::Tensor & self, int64_t dim, bool descending) {
  18100. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18101. auto maybe_layer = maybeCurrentDynamicLayer();
  18102. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18103. int64_t cur_level = maybe_layer->layerId();
  18104. if (!isBatchedAtLevel(self, cur_level)) {
  18105. return at::_ops::argsort::call(self, dim, descending);
  18106. }
  18107. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  18108. auto results = batch_rule(self_value, self_bdim, dim, descending);
  18109. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  18110. }
  18111. template <typename batch_rule_t, batch_rule_t batch_rule>
  18112. at::Tensor argsort_stable_generated_plumbing(const at::Tensor & self, bool stable, int64_t dim, bool descending) {
  18113. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18114. auto maybe_layer = maybeCurrentDynamicLayer();
  18115. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18116. int64_t cur_level = maybe_layer->layerId();
  18117. if (!isBatchedAtLevel(self, cur_level)) {
  18118. return at::_ops::argsort_stable::call(self, stable, dim, descending);
  18119. }
  18120. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  18121. auto results = batch_rule(self_value, self_bdim, stable, dim, descending);
  18122. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  18123. }
  18124. template <typename batch_rule_t, batch_rule_t batch_rule>
  18125. at::Tensor argsort_dimname_generated_plumbing(const at::Tensor & self, at::Dimname dim, bool descending) {
  18126. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18127. auto maybe_layer = maybeCurrentDynamicLayer();
  18128. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18129. int64_t cur_level = maybe_layer->layerId();
  18130. if (!isBatchedAtLevel(self, cur_level)) {
  18131. return at::_ops::argsort_dimname::call(self, dim, descending);
  18132. }
  18133. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  18134. auto results = batch_rule(self_value, self_bdim, dim, descending);
  18135. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  18136. }
  18137. template <typename batch_rule_t, batch_rule_t batch_rule>
  18138. ::std::tuple<at::Tensor,at::Tensor> topk_generated_plumbing(const at::Tensor & self, c10::SymInt k, int64_t dim, bool largest, bool sorted) {
  18139. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18140. auto maybe_layer = maybeCurrentDynamicLayer();
  18141. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18142. int64_t cur_level = maybe_layer->layerId();
  18143. if (!isBatchedAtLevel(self, cur_level)) {
  18144. return at::_ops::topk::call(self, k, dim, largest, sorted);
  18145. }
  18146. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  18147. auto results = batch_rule(self_value, self_bdim, k, dim, largest, sorted);
  18148. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  18149. }
  18150. template <typename batch_rule_t, batch_rule_t batch_rule>
  18151. at::Tensor all_generated_plumbing(const at::Tensor & self) {
  18152. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18153. auto maybe_layer = maybeCurrentDynamicLayer();
  18154. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18155. int64_t cur_level = maybe_layer->layerId();
  18156. if (!isBatchedAtLevel(self, cur_level)) {
  18157. return at::_ops::all::call(self);
  18158. }
  18159. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  18160. auto results = batch_rule(self_value, self_bdim);
  18161. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  18162. }
  18163. template <typename batch_rule_t, batch_rule_t batch_rule>
  18164. at::Tensor any_generated_plumbing(const at::Tensor & self) {
  18165. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18166. auto maybe_layer = maybeCurrentDynamicLayer();
  18167. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18168. int64_t cur_level = maybe_layer->layerId();
  18169. if (!isBatchedAtLevel(self, cur_level)) {
  18170. return at::_ops::any::call(self);
  18171. }
  18172. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  18173. auto results = batch_rule(self_value, self_bdim);
  18174. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  18175. }
  18176. template <typename batch_rule_t, batch_rule_t batch_rule>
  18177. at::Tensor renorm_generated_plumbing(const at::Tensor & self, const at::Scalar & p, int64_t dim, const at::Scalar & maxnorm) {
  18178. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18179. auto maybe_layer = maybeCurrentDynamicLayer();
  18180. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18181. int64_t cur_level = maybe_layer->layerId();
  18182. if (!isBatchedAtLevel(self, cur_level)) {
  18183. return at::_ops::renorm::call(self, p, dim, maxnorm);
  18184. }
  18185. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  18186. auto results = batch_rule(self_value, self_bdim, p, dim, maxnorm);
  18187. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  18188. }
  18189. template <typename batch_rule_t, batch_rule_t batch_rule>
  18190. at::Tensor & renorm__generated_plumbing(at::Tensor & self, const at::Scalar & p, int64_t dim, const at::Scalar & maxnorm) {
  18191. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18192. auto maybe_layer = maybeCurrentDynamicLayer();
  18193. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  18194. int64_t cur_level = maybe_layer->layerId();
  18195. if (!isBatchedAtLevel(self, cur_level)) {
  18196. return at::_ops::renorm_::call(self, p, dim, maxnorm);
  18197. }
  18198. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  18199. batch_rule(self_value, self_bdim, p, dim, maxnorm);
  18200. return self;
  18201. }
  18202. template <typename batch_rule_t, batch_rule_t batch_rule>
  18203. at::Tensor unfold_generated_plumbing(const at::Tensor & self, int64_t dimension, int64_t size, int64_t step) {
  18204. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18205. auto maybe_layer = maybeCurrentDynamicLayer();
  18206. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18207. int64_t cur_level = maybe_layer->layerId();
  18208. if (!isBatchedAtLevel(self, cur_level)) {
  18209. return at::_ops::unfold::call(self, dimension, size, step);
  18210. }
  18211. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  18212. auto results = batch_rule(self_value, self_bdim, dimension, size, step);
  18213. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  18214. }
  18215. template <typename batch_rule_t, batch_rule_t batch_rule>
  18216. at::Tensor unfold_backward_generated_plumbing(const at::Tensor & grad_in, c10::SymIntArrayRef input_sizes, int64_t dim, int64_t size, int64_t step) {
  18217. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18218. auto maybe_layer = maybeCurrentDynamicLayer();
  18219. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18220. int64_t cur_level = maybe_layer->layerId();
  18221. if (!isBatchedAtLevel(grad_in, cur_level)) {
  18222. return at::_ops::unfold_backward::call(grad_in, input_sizes, dim, size, step);
  18223. }
  18224. auto [grad_in_value, grad_in_bdim] = unwrapTensorAtLevel(grad_in, cur_level);
  18225. auto results = batch_rule(grad_in_value, grad_in_bdim, input_sizes, dim, size, step);
  18226. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  18227. }
  18228. template <typename batch_rule_t, batch_rule_t batch_rule>
  18229. at::Tensor pow_Tensor_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & exponent) {
  18230. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18231. auto maybe_layer = maybeCurrentDynamicLayer();
  18232. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18233. int64_t cur_level = maybe_layer->layerId();
  18234. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(exponent, cur_level)) {
  18235. return at::_ops::pow_Tensor_Tensor::call(self, exponent);
  18236. }
  18237. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  18238. auto [exponent_value, exponent_bdim] = unwrapTensorAtLevel(exponent, cur_level);
  18239. auto results = batch_rule(self_value, self_bdim, exponent_value, exponent_bdim);
  18240. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  18241. }
  18242. template <typename batch_rule_t, batch_rule_t batch_rule>
  18243. at::Tensor pow_Scalar_generated_plumbing(const at::Scalar & self, const at::Tensor & exponent) {
  18244. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18245. auto maybe_layer = maybeCurrentDynamicLayer();
  18246. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18247. int64_t cur_level = maybe_layer->layerId();
  18248. if (!isBatchedAtLevel(exponent, cur_level)) {
  18249. return at::_ops::pow_Scalar::call(self, exponent);
  18250. }
  18251. auto [exponent_value, exponent_bdim] = unwrapTensorAtLevel(exponent, cur_level);
  18252. auto results = batch_rule(self, exponent_value, exponent_bdim);
  18253. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  18254. }
  18255. template <typename batch_rule_t, batch_rule_t batch_rule>
  18256. at::Tensor pow_Tensor_Scalar_generated_plumbing(const at::Tensor & self, const at::Scalar & exponent) {
  18257. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18258. auto maybe_layer = maybeCurrentDynamicLayer();
  18259. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18260. int64_t cur_level = maybe_layer->layerId();
  18261. if (!isBatchedAtLevel(self, cur_level)) {
  18262. return at::_ops::pow_Tensor_Scalar::call(self, exponent);
  18263. }
  18264. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  18265. auto results = batch_rule(self_value, self_bdim, exponent);
  18266. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  18267. }
  18268. template <typename batch_rule_t, batch_rule_t batch_rule>
  18269. at::Tensor & pow__Scalar_generated_plumbing(at::Tensor & self, const at::Scalar & exponent) {
  18270. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18271. auto maybe_layer = maybeCurrentDynamicLayer();
  18272. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  18273. int64_t cur_level = maybe_layer->layerId();
  18274. if (!isBatchedAtLevel(self, cur_level)) {
  18275. return at::_ops::pow__Scalar::call(self, exponent);
  18276. }
  18277. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  18278. batch_rule(self_value, self_bdim, exponent);
  18279. return self;
  18280. }
  18281. template <typename batch_rule_t, batch_rule_t batch_rule>
  18282. at::Tensor & pow__Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & exponent) {
  18283. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18284. auto maybe_layer = maybeCurrentDynamicLayer();
  18285. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  18286. int64_t cur_level = maybe_layer->layerId();
  18287. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(exponent, cur_level)) {
  18288. return at::_ops::pow__Tensor::call(self, exponent);
  18289. }
  18290. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  18291. auto [exponent_value, exponent_bdim] = unwrapTensorAtLevel(exponent, cur_level);
  18292. batch_rule(self_value, self_bdim, exponent_value, exponent_bdim);
  18293. return self;
  18294. }
  18295. template <typename batch_rule_t, batch_rule_t batch_rule>
  18296. at::Tensor float_power_Tensor_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & exponent) {
  18297. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18298. auto maybe_layer = maybeCurrentDynamicLayer();
  18299. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18300. int64_t cur_level = maybe_layer->layerId();
  18301. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(exponent, cur_level)) {
  18302. return at::_ops::float_power_Tensor_Tensor::call(self, exponent);
  18303. }
  18304. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  18305. auto [exponent_value, exponent_bdim] = unwrapTensorAtLevel(exponent, cur_level);
  18306. auto results = batch_rule(self_value, self_bdim, exponent_value, exponent_bdim);
  18307. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  18308. }
  18309. template <typename batch_rule_t, batch_rule_t batch_rule>
  18310. at::Tensor float_power_Scalar_generated_plumbing(const at::Scalar & self, const at::Tensor & exponent) {
  18311. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18312. auto maybe_layer = maybeCurrentDynamicLayer();
  18313. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18314. int64_t cur_level = maybe_layer->layerId();
  18315. if (!isBatchedAtLevel(exponent, cur_level)) {
  18316. return at::_ops::float_power_Scalar::call(self, exponent);
  18317. }
  18318. auto [exponent_value, exponent_bdim] = unwrapTensorAtLevel(exponent, cur_level);
  18319. auto results = batch_rule(self, exponent_value, exponent_bdim);
  18320. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  18321. }
  18322. template <typename batch_rule_t, batch_rule_t batch_rule>
  18323. at::Tensor float_power_Tensor_Scalar_generated_plumbing(const at::Tensor & self, const at::Scalar & exponent) {
  18324. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18325. auto maybe_layer = maybeCurrentDynamicLayer();
  18326. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18327. int64_t cur_level = maybe_layer->layerId();
  18328. if (!isBatchedAtLevel(self, cur_level)) {
  18329. return at::_ops::float_power_Tensor_Scalar::call(self, exponent);
  18330. }
  18331. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  18332. auto results = batch_rule(self_value, self_bdim, exponent);
  18333. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  18334. }
  18335. template <typename batch_rule_t, batch_rule_t batch_rule>
  18336. at::Tensor & float_power__Scalar_generated_plumbing(at::Tensor & self, const at::Scalar & exponent) {
  18337. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18338. auto maybe_layer = maybeCurrentDynamicLayer();
  18339. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  18340. int64_t cur_level = maybe_layer->layerId();
  18341. if (!isBatchedAtLevel(self, cur_level)) {
  18342. return at::_ops::float_power__Scalar::call(self, exponent);
  18343. }
  18344. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  18345. batch_rule(self_value, self_bdim, exponent);
  18346. return self;
  18347. }
  18348. template <typename batch_rule_t, batch_rule_t batch_rule>
  18349. at::Tensor & float_power__Tensor_generated_plumbing(at::Tensor & self, const at::Tensor & exponent) {
  18350. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18351. auto maybe_layer = maybeCurrentDynamicLayer();
  18352. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  18353. int64_t cur_level = maybe_layer->layerId();
  18354. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(exponent, cur_level)) {
  18355. return at::_ops::float_power__Tensor::call(self, exponent);
  18356. }
  18357. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  18358. auto [exponent_value, exponent_bdim] = unwrapTensorAtLevel(exponent, cur_level);
  18359. batch_rule(self_value, self_bdim, exponent_value, exponent_bdim);
  18360. return self;
  18361. }
  18362. template <typename batch_rule_t, batch_rule_t batch_rule>
  18363. at::Tensor & normal__generated_plumbing(at::Tensor & self, double mean, double std, ::std::optional<at::Generator> generator) {
  18364. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18365. auto maybe_layer = maybeCurrentDynamicLayer();
  18366. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  18367. int64_t cur_level = maybe_layer->layerId();
  18368. if (!isBatchedAtLevel(self, cur_level)) {
  18369. return at::_ops::normal_::call(self, mean, std, generator);
  18370. }
  18371. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  18372. batch_rule(self_value, self_bdim, mean, std, generator);
  18373. return self;
  18374. }
  18375. template <typename batch_rule_t, batch_rule_t batch_rule>
  18376. at::Tensor normal_functional_generated_plumbing(const at::Tensor & self, double mean, double std, ::std::optional<at::Generator> generator) {
  18377. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18378. auto maybe_layer = maybeCurrentDynamicLayer();
  18379. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18380. int64_t cur_level = maybe_layer->layerId();
  18381. if (!isBatchedAtLevel(self, cur_level)) {
  18382. return at::_ops::normal_functional::call(self, mean, std, generator);
  18383. }
  18384. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  18385. auto results = batch_rule(self_value, self_bdim, mean, std, generator);
  18386. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  18387. }
  18388. template <typename batch_rule_t, batch_rule_t batch_rule>
  18389. at::Tensor normal_Tensor_float_generated_plumbing(const at::Tensor & mean, double std, ::std::optional<at::Generator> generator) {
  18390. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18391. auto maybe_layer = maybeCurrentDynamicLayer();
  18392. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18393. int64_t cur_level = maybe_layer->layerId();
  18394. if (!isBatchedAtLevel(mean, cur_level)) {
  18395. return at::_ops::normal_Tensor_float::call(mean, std, generator);
  18396. }
  18397. auto [mean_value, mean_bdim] = unwrapTensorAtLevel(mean, cur_level);
  18398. auto results = batch_rule(mean_value, mean_bdim, std, generator);
  18399. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  18400. }
  18401. template <typename batch_rule_t, batch_rule_t batch_rule>
  18402. at::Tensor normal_float_Tensor_generated_plumbing(double mean, const at::Tensor & std, ::std::optional<at::Generator> generator) {
  18403. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18404. auto maybe_layer = maybeCurrentDynamicLayer();
  18405. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18406. int64_t cur_level = maybe_layer->layerId();
  18407. if (!isBatchedAtLevel(std, cur_level)) {
  18408. return at::_ops::normal_float_Tensor::call(mean, std, generator);
  18409. }
  18410. auto [std_value, std_bdim] = unwrapTensorAtLevel(std, cur_level);
  18411. auto results = batch_rule(mean, std_value, std_bdim, generator);
  18412. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  18413. }
  18414. template <typename batch_rule_t, batch_rule_t batch_rule>
  18415. at::Tensor normal_Tensor_Tensor_generated_plumbing(const at::Tensor & mean, const at::Tensor & std, ::std::optional<at::Generator> generator) {
  18416. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18417. auto maybe_layer = maybeCurrentDynamicLayer();
  18418. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18419. int64_t cur_level = maybe_layer->layerId();
  18420. if (!isBatchedAtLevel(mean, cur_level) && !isBatchedAtLevel(std, cur_level)) {
  18421. return at::_ops::normal_Tensor_Tensor::call(mean, std, generator);
  18422. }
  18423. auto [mean_value, mean_bdim] = unwrapTensorAtLevel(mean, cur_level);
  18424. auto [std_value, std_bdim] = unwrapTensorAtLevel(std, cur_level);
  18425. auto results = batch_rule(mean_value, mean_bdim, std_value, std_bdim, generator);
  18426. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  18427. }
  18428. template <typename batch_rule_t, batch_rule_t batch_rule>
  18429. at::Tensor alias_generated_plumbing(const at::Tensor & self) {
  18430. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18431. auto maybe_layer = maybeCurrentDynamicLayer();
  18432. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18433. int64_t cur_level = maybe_layer->layerId();
  18434. if (!isBatchedAtLevel(self, cur_level)) {
  18435. return at::_ops::alias::call(self);
  18436. }
  18437. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  18438. auto results = batch_rule(self_value, self_bdim);
  18439. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  18440. }
  18441. template <typename batch_rule_t, batch_rule_t batch_rule>
  18442. void _amp_foreach_non_finite_check_and_unscale__generated_plumbing(at::TensorList self, at::Tensor & found_inf, const at::Tensor & inv_scale) {
  18443. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18444. auto maybe_layer = maybeCurrentDynamicLayer();
  18445. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  18446. int64_t cur_level = maybe_layer->layerId();
  18447. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(found_inf, cur_level) && !isBatchedAtLevel(inv_scale, cur_level)) {
  18448. return at::_ops::_amp_foreach_non_finite_check_and_unscale_::call(self, found_inf, inv_scale);
  18449. }
  18450. auto [found_inf_value, found_inf_bdim] = unwrapTensorAtLevel(found_inf, cur_level);
  18451. auto [inv_scale_value, inv_scale_bdim] = unwrapTensorAtLevel(inv_scale, cur_level);
  18452. batch_rule(self, found_inf_value, found_inf_bdim, inv_scale_value, inv_scale_bdim);
  18453. }
  18454. template <typename batch_rule_t, batch_rule_t batch_rule>
  18455. ::std::vector<at::Tensor> _foreach_add_Scalar_generated_plumbing(at::TensorList self, const at::Scalar & scalar) {
  18456. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18457. auto maybe_layer = maybeCurrentDynamicLayer();
  18458. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18459. int64_t cur_level = maybe_layer->layerId();
  18460. if (!isBatchedAtLevel(self, cur_level)) {
  18461. return at::_ops::_foreach_add_Scalar::call(self, scalar);
  18462. }
  18463. auto results = batch_rule(self, scalar);
  18464. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  18465. }
  18466. template <typename batch_rule_t, batch_rule_t batch_rule>
  18467. void _foreach_add__Scalar_generated_plumbing(at::TensorList self, const at::Scalar & scalar) {
  18468. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18469. auto maybe_layer = maybeCurrentDynamicLayer();
  18470. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  18471. int64_t cur_level = maybe_layer->layerId();
  18472. if (!isBatchedAtLevel(self, cur_level)) {
  18473. return at::_ops::_foreach_add__Scalar::call(self, scalar);
  18474. }
  18475. batch_rule(self, scalar);
  18476. }
  18477. template <typename batch_rule_t, batch_rule_t batch_rule>
  18478. ::std::vector<at::Tensor> _foreach_add_List_generated_plumbing(at::TensorList self, at::TensorList other, const at::Scalar & alpha) {
  18479. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18480. auto maybe_layer = maybeCurrentDynamicLayer();
  18481. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18482. int64_t cur_level = maybe_layer->layerId();
  18483. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  18484. return at::_ops::_foreach_add_List::call(self, other, alpha);
  18485. }
  18486. auto results = batch_rule(self, other, alpha);
  18487. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  18488. }
  18489. template <typename batch_rule_t, batch_rule_t batch_rule>
  18490. void _foreach_add__List_generated_plumbing(at::TensorList self, at::TensorList other, const at::Scalar & alpha) {
  18491. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18492. auto maybe_layer = maybeCurrentDynamicLayer();
  18493. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  18494. int64_t cur_level = maybe_layer->layerId();
  18495. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  18496. return at::_ops::_foreach_add__List::call(self, other, alpha);
  18497. }
  18498. batch_rule(self, other, alpha);
  18499. }
  18500. template <typename batch_rule_t, batch_rule_t batch_rule>
  18501. ::std::vector<at::Tensor> _foreach_add_ScalarList_generated_plumbing(at::TensorList self, at::ArrayRef<at::Scalar> scalars) {
  18502. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18503. auto maybe_layer = maybeCurrentDynamicLayer();
  18504. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18505. int64_t cur_level = maybe_layer->layerId();
  18506. if (!isBatchedAtLevel(self, cur_level)) {
  18507. return at::_ops::_foreach_add_ScalarList::call(self, scalars);
  18508. }
  18509. auto results = batch_rule(self, scalars);
  18510. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  18511. }
  18512. template <typename batch_rule_t, batch_rule_t batch_rule>
  18513. void _foreach_add__ScalarList_generated_plumbing(at::TensorList self, at::ArrayRef<at::Scalar> scalars) {
  18514. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18515. auto maybe_layer = maybeCurrentDynamicLayer();
  18516. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  18517. int64_t cur_level = maybe_layer->layerId();
  18518. if (!isBatchedAtLevel(self, cur_level)) {
  18519. return at::_ops::_foreach_add__ScalarList::call(self, scalars);
  18520. }
  18521. batch_rule(self, scalars);
  18522. }
  18523. template <typename batch_rule_t, batch_rule_t batch_rule>
  18524. ::std::vector<at::Tensor> _foreach_add_Tensor_generated_plumbing(at::TensorList self, const at::Tensor & other, const at::Scalar & alpha) {
  18525. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18526. auto maybe_layer = maybeCurrentDynamicLayer();
  18527. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18528. int64_t cur_level = maybe_layer->layerId();
  18529. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  18530. return at::_ops::_foreach_add_Tensor::call(self, other, alpha);
  18531. }
  18532. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  18533. auto results = batch_rule(self, other_value, other_bdim, alpha);
  18534. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  18535. }
  18536. template <typename batch_rule_t, batch_rule_t batch_rule>
  18537. void _foreach_add__Tensor_generated_plumbing(at::TensorList self, const at::Tensor & other, const at::Scalar & alpha) {
  18538. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18539. auto maybe_layer = maybeCurrentDynamicLayer();
  18540. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  18541. int64_t cur_level = maybe_layer->layerId();
  18542. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  18543. return at::_ops::_foreach_add__Tensor::call(self, other, alpha);
  18544. }
  18545. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  18546. batch_rule(self, other_value, other_bdim, alpha);
  18547. }
  18548. template <typename batch_rule_t, batch_rule_t batch_rule>
  18549. ::std::vector<at::Tensor> _foreach_sub_Scalar_generated_plumbing(at::TensorList self, const at::Scalar & scalar) {
  18550. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18551. auto maybe_layer = maybeCurrentDynamicLayer();
  18552. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18553. int64_t cur_level = maybe_layer->layerId();
  18554. if (!isBatchedAtLevel(self, cur_level)) {
  18555. return at::_ops::_foreach_sub_Scalar::call(self, scalar);
  18556. }
  18557. auto results = batch_rule(self, scalar);
  18558. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  18559. }
  18560. template <typename batch_rule_t, batch_rule_t batch_rule>
  18561. void _foreach_sub__Scalar_generated_plumbing(at::TensorList self, const at::Scalar & scalar) {
  18562. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18563. auto maybe_layer = maybeCurrentDynamicLayer();
  18564. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  18565. int64_t cur_level = maybe_layer->layerId();
  18566. if (!isBatchedAtLevel(self, cur_level)) {
  18567. return at::_ops::_foreach_sub__Scalar::call(self, scalar);
  18568. }
  18569. batch_rule(self, scalar);
  18570. }
  18571. template <typename batch_rule_t, batch_rule_t batch_rule>
  18572. ::std::vector<at::Tensor> _foreach_sub_List_generated_plumbing(at::TensorList self, at::TensorList other, const at::Scalar & alpha) {
  18573. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18574. auto maybe_layer = maybeCurrentDynamicLayer();
  18575. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18576. int64_t cur_level = maybe_layer->layerId();
  18577. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  18578. return at::_ops::_foreach_sub_List::call(self, other, alpha);
  18579. }
  18580. auto results = batch_rule(self, other, alpha);
  18581. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  18582. }
  18583. template <typename batch_rule_t, batch_rule_t batch_rule>
  18584. void _foreach_sub__List_generated_plumbing(at::TensorList self, at::TensorList other, const at::Scalar & alpha) {
  18585. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18586. auto maybe_layer = maybeCurrentDynamicLayer();
  18587. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  18588. int64_t cur_level = maybe_layer->layerId();
  18589. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  18590. return at::_ops::_foreach_sub__List::call(self, other, alpha);
  18591. }
  18592. batch_rule(self, other, alpha);
  18593. }
  18594. template <typename batch_rule_t, batch_rule_t batch_rule>
  18595. ::std::vector<at::Tensor> _foreach_sub_ScalarList_generated_plumbing(at::TensorList self, at::ArrayRef<at::Scalar> scalars) {
  18596. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18597. auto maybe_layer = maybeCurrentDynamicLayer();
  18598. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18599. int64_t cur_level = maybe_layer->layerId();
  18600. if (!isBatchedAtLevel(self, cur_level)) {
  18601. return at::_ops::_foreach_sub_ScalarList::call(self, scalars);
  18602. }
  18603. auto results = batch_rule(self, scalars);
  18604. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  18605. }
  18606. template <typename batch_rule_t, batch_rule_t batch_rule>
  18607. void _foreach_sub__ScalarList_generated_plumbing(at::TensorList self, at::ArrayRef<at::Scalar> scalars) {
  18608. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18609. auto maybe_layer = maybeCurrentDynamicLayer();
  18610. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  18611. int64_t cur_level = maybe_layer->layerId();
  18612. if (!isBatchedAtLevel(self, cur_level)) {
  18613. return at::_ops::_foreach_sub__ScalarList::call(self, scalars);
  18614. }
  18615. batch_rule(self, scalars);
  18616. }
  18617. template <typename batch_rule_t, batch_rule_t batch_rule>
  18618. ::std::vector<at::Tensor> _foreach_mul_Scalar_generated_plumbing(at::TensorList self, const at::Scalar & scalar) {
  18619. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18620. auto maybe_layer = maybeCurrentDynamicLayer();
  18621. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18622. int64_t cur_level = maybe_layer->layerId();
  18623. if (!isBatchedAtLevel(self, cur_level)) {
  18624. return at::_ops::_foreach_mul_Scalar::call(self, scalar);
  18625. }
  18626. auto results = batch_rule(self, scalar);
  18627. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  18628. }
  18629. template <typename batch_rule_t, batch_rule_t batch_rule>
  18630. void _foreach_mul__Scalar_generated_plumbing(at::TensorList self, const at::Scalar & scalar) {
  18631. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18632. auto maybe_layer = maybeCurrentDynamicLayer();
  18633. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  18634. int64_t cur_level = maybe_layer->layerId();
  18635. if (!isBatchedAtLevel(self, cur_level)) {
  18636. return at::_ops::_foreach_mul__Scalar::call(self, scalar);
  18637. }
  18638. batch_rule(self, scalar);
  18639. }
  18640. template <typename batch_rule_t, batch_rule_t batch_rule>
  18641. ::std::vector<at::Tensor> _foreach_mul_List_generated_plumbing(at::TensorList self, at::TensorList other) {
  18642. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18643. auto maybe_layer = maybeCurrentDynamicLayer();
  18644. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18645. int64_t cur_level = maybe_layer->layerId();
  18646. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  18647. return at::_ops::_foreach_mul_List::call(self, other);
  18648. }
  18649. auto results = batch_rule(self, other);
  18650. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  18651. }
  18652. template <typename batch_rule_t, batch_rule_t batch_rule>
  18653. void _foreach_mul__List_generated_plumbing(at::TensorList self, at::TensorList other) {
  18654. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18655. auto maybe_layer = maybeCurrentDynamicLayer();
  18656. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  18657. int64_t cur_level = maybe_layer->layerId();
  18658. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  18659. return at::_ops::_foreach_mul__List::call(self, other);
  18660. }
  18661. batch_rule(self, other);
  18662. }
  18663. template <typename batch_rule_t, batch_rule_t batch_rule>
  18664. ::std::vector<at::Tensor> _foreach_mul_ScalarList_generated_plumbing(at::TensorList self, at::ArrayRef<at::Scalar> scalars) {
  18665. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18666. auto maybe_layer = maybeCurrentDynamicLayer();
  18667. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18668. int64_t cur_level = maybe_layer->layerId();
  18669. if (!isBatchedAtLevel(self, cur_level)) {
  18670. return at::_ops::_foreach_mul_ScalarList::call(self, scalars);
  18671. }
  18672. auto results = batch_rule(self, scalars);
  18673. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  18674. }
  18675. template <typename batch_rule_t, batch_rule_t batch_rule>
  18676. void _foreach_mul__ScalarList_generated_plumbing(at::TensorList self, at::ArrayRef<at::Scalar> scalars) {
  18677. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18678. auto maybe_layer = maybeCurrentDynamicLayer();
  18679. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  18680. int64_t cur_level = maybe_layer->layerId();
  18681. if (!isBatchedAtLevel(self, cur_level)) {
  18682. return at::_ops::_foreach_mul__ScalarList::call(self, scalars);
  18683. }
  18684. batch_rule(self, scalars);
  18685. }
  18686. template <typename batch_rule_t, batch_rule_t batch_rule>
  18687. ::std::vector<at::Tensor> _foreach_mul_Tensor_generated_plumbing(at::TensorList self, const at::Tensor & other) {
  18688. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18689. auto maybe_layer = maybeCurrentDynamicLayer();
  18690. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18691. int64_t cur_level = maybe_layer->layerId();
  18692. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  18693. return at::_ops::_foreach_mul_Tensor::call(self, other);
  18694. }
  18695. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  18696. auto results = batch_rule(self, other_value, other_bdim);
  18697. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  18698. }
  18699. template <typename batch_rule_t, batch_rule_t batch_rule>
  18700. void _foreach_mul__Tensor_generated_plumbing(at::TensorList self, const at::Tensor & other) {
  18701. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18702. auto maybe_layer = maybeCurrentDynamicLayer();
  18703. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  18704. int64_t cur_level = maybe_layer->layerId();
  18705. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  18706. return at::_ops::_foreach_mul__Tensor::call(self, other);
  18707. }
  18708. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  18709. batch_rule(self, other_value, other_bdim);
  18710. }
  18711. template <typename batch_rule_t, batch_rule_t batch_rule>
  18712. ::std::vector<at::Tensor> _foreach_div_Scalar_generated_plumbing(at::TensorList self, const at::Scalar & scalar) {
  18713. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18714. auto maybe_layer = maybeCurrentDynamicLayer();
  18715. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18716. int64_t cur_level = maybe_layer->layerId();
  18717. if (!isBatchedAtLevel(self, cur_level)) {
  18718. return at::_ops::_foreach_div_Scalar::call(self, scalar);
  18719. }
  18720. auto results = batch_rule(self, scalar);
  18721. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  18722. }
  18723. template <typename batch_rule_t, batch_rule_t batch_rule>
  18724. void _foreach_div__Scalar_generated_plumbing(at::TensorList self, const at::Scalar & scalar) {
  18725. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18726. auto maybe_layer = maybeCurrentDynamicLayer();
  18727. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  18728. int64_t cur_level = maybe_layer->layerId();
  18729. if (!isBatchedAtLevel(self, cur_level)) {
  18730. return at::_ops::_foreach_div__Scalar::call(self, scalar);
  18731. }
  18732. batch_rule(self, scalar);
  18733. }
  18734. template <typename batch_rule_t, batch_rule_t batch_rule>
  18735. ::std::vector<at::Tensor> _foreach_div_List_generated_plumbing(at::TensorList self, at::TensorList other) {
  18736. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18737. auto maybe_layer = maybeCurrentDynamicLayer();
  18738. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18739. int64_t cur_level = maybe_layer->layerId();
  18740. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  18741. return at::_ops::_foreach_div_List::call(self, other);
  18742. }
  18743. auto results = batch_rule(self, other);
  18744. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  18745. }
  18746. template <typename batch_rule_t, batch_rule_t batch_rule>
  18747. void _foreach_div__List_generated_plumbing(at::TensorList self, at::TensorList other) {
  18748. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18749. auto maybe_layer = maybeCurrentDynamicLayer();
  18750. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  18751. int64_t cur_level = maybe_layer->layerId();
  18752. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  18753. return at::_ops::_foreach_div__List::call(self, other);
  18754. }
  18755. batch_rule(self, other);
  18756. }
  18757. template <typename batch_rule_t, batch_rule_t batch_rule>
  18758. ::std::vector<at::Tensor> _foreach_div_ScalarList_generated_plumbing(at::TensorList self, at::ArrayRef<at::Scalar> scalars) {
  18759. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18760. auto maybe_layer = maybeCurrentDynamicLayer();
  18761. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18762. int64_t cur_level = maybe_layer->layerId();
  18763. if (!isBatchedAtLevel(self, cur_level)) {
  18764. return at::_ops::_foreach_div_ScalarList::call(self, scalars);
  18765. }
  18766. auto results = batch_rule(self, scalars);
  18767. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  18768. }
  18769. template <typename batch_rule_t, batch_rule_t batch_rule>
  18770. void _foreach_div__ScalarList_generated_plumbing(at::TensorList self, at::ArrayRef<at::Scalar> scalars) {
  18771. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18772. auto maybe_layer = maybeCurrentDynamicLayer();
  18773. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  18774. int64_t cur_level = maybe_layer->layerId();
  18775. if (!isBatchedAtLevel(self, cur_level)) {
  18776. return at::_ops::_foreach_div__ScalarList::call(self, scalars);
  18777. }
  18778. batch_rule(self, scalars);
  18779. }
  18780. template <typename batch_rule_t, batch_rule_t batch_rule>
  18781. ::std::vector<at::Tensor> _foreach_div_Tensor_generated_plumbing(at::TensorList self, const at::Tensor & other) {
  18782. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18783. auto maybe_layer = maybeCurrentDynamicLayer();
  18784. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18785. int64_t cur_level = maybe_layer->layerId();
  18786. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  18787. return at::_ops::_foreach_div_Tensor::call(self, other);
  18788. }
  18789. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  18790. auto results = batch_rule(self, other_value, other_bdim);
  18791. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  18792. }
  18793. template <typename batch_rule_t, batch_rule_t batch_rule>
  18794. void _foreach_div__Tensor_generated_plumbing(at::TensorList self, const at::Tensor & other) {
  18795. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18796. auto maybe_layer = maybeCurrentDynamicLayer();
  18797. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  18798. int64_t cur_level = maybe_layer->layerId();
  18799. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  18800. return at::_ops::_foreach_div__Tensor::call(self, other);
  18801. }
  18802. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  18803. batch_rule(self, other_value, other_bdim);
  18804. }
  18805. template <typename batch_rule_t, batch_rule_t batch_rule>
  18806. ::std::vector<at::Tensor> _foreach_clamp_max_Scalar_generated_plumbing(at::TensorList self, const at::Scalar & scalar) {
  18807. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18808. auto maybe_layer = maybeCurrentDynamicLayer();
  18809. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18810. int64_t cur_level = maybe_layer->layerId();
  18811. if (!isBatchedAtLevel(self, cur_level)) {
  18812. return at::_ops::_foreach_clamp_max_Scalar::call(self, scalar);
  18813. }
  18814. auto results = batch_rule(self, scalar);
  18815. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  18816. }
  18817. template <typename batch_rule_t, batch_rule_t batch_rule>
  18818. void _foreach_clamp_max__Scalar_generated_plumbing(at::TensorList self, const at::Scalar & scalar) {
  18819. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18820. auto maybe_layer = maybeCurrentDynamicLayer();
  18821. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  18822. int64_t cur_level = maybe_layer->layerId();
  18823. if (!isBatchedAtLevel(self, cur_level)) {
  18824. return at::_ops::_foreach_clamp_max__Scalar::call(self, scalar);
  18825. }
  18826. batch_rule(self, scalar);
  18827. }
  18828. template <typename batch_rule_t, batch_rule_t batch_rule>
  18829. ::std::vector<at::Tensor> _foreach_clamp_max_List_generated_plumbing(at::TensorList self, at::TensorList other) {
  18830. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18831. auto maybe_layer = maybeCurrentDynamicLayer();
  18832. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18833. int64_t cur_level = maybe_layer->layerId();
  18834. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  18835. return at::_ops::_foreach_clamp_max_List::call(self, other);
  18836. }
  18837. auto results = batch_rule(self, other);
  18838. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  18839. }
  18840. template <typename batch_rule_t, batch_rule_t batch_rule>
  18841. void _foreach_clamp_max__List_generated_plumbing(at::TensorList self, at::TensorList other) {
  18842. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18843. auto maybe_layer = maybeCurrentDynamicLayer();
  18844. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  18845. int64_t cur_level = maybe_layer->layerId();
  18846. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  18847. return at::_ops::_foreach_clamp_max__List::call(self, other);
  18848. }
  18849. batch_rule(self, other);
  18850. }
  18851. template <typename batch_rule_t, batch_rule_t batch_rule>
  18852. ::std::vector<at::Tensor> _foreach_clamp_max_ScalarList_generated_plumbing(at::TensorList self, at::ArrayRef<at::Scalar> scalars) {
  18853. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18854. auto maybe_layer = maybeCurrentDynamicLayer();
  18855. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18856. int64_t cur_level = maybe_layer->layerId();
  18857. if (!isBatchedAtLevel(self, cur_level)) {
  18858. return at::_ops::_foreach_clamp_max_ScalarList::call(self, scalars);
  18859. }
  18860. auto results = batch_rule(self, scalars);
  18861. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  18862. }
  18863. template <typename batch_rule_t, batch_rule_t batch_rule>
  18864. void _foreach_clamp_max__ScalarList_generated_plumbing(at::TensorList self, at::ArrayRef<at::Scalar> scalars) {
  18865. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18866. auto maybe_layer = maybeCurrentDynamicLayer();
  18867. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  18868. int64_t cur_level = maybe_layer->layerId();
  18869. if (!isBatchedAtLevel(self, cur_level)) {
  18870. return at::_ops::_foreach_clamp_max__ScalarList::call(self, scalars);
  18871. }
  18872. batch_rule(self, scalars);
  18873. }
  18874. template <typename batch_rule_t, batch_rule_t batch_rule>
  18875. ::std::vector<at::Tensor> _foreach_clamp_min_Scalar_generated_plumbing(at::TensorList self, const at::Scalar & scalar) {
  18876. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18877. auto maybe_layer = maybeCurrentDynamicLayer();
  18878. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18879. int64_t cur_level = maybe_layer->layerId();
  18880. if (!isBatchedAtLevel(self, cur_level)) {
  18881. return at::_ops::_foreach_clamp_min_Scalar::call(self, scalar);
  18882. }
  18883. auto results = batch_rule(self, scalar);
  18884. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  18885. }
  18886. template <typename batch_rule_t, batch_rule_t batch_rule>
  18887. void _foreach_clamp_min__Scalar_generated_plumbing(at::TensorList self, const at::Scalar & scalar) {
  18888. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18889. auto maybe_layer = maybeCurrentDynamicLayer();
  18890. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  18891. int64_t cur_level = maybe_layer->layerId();
  18892. if (!isBatchedAtLevel(self, cur_level)) {
  18893. return at::_ops::_foreach_clamp_min__Scalar::call(self, scalar);
  18894. }
  18895. batch_rule(self, scalar);
  18896. }
  18897. template <typename batch_rule_t, batch_rule_t batch_rule>
  18898. ::std::vector<at::Tensor> _foreach_clamp_min_List_generated_plumbing(at::TensorList self, at::TensorList other) {
  18899. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18900. auto maybe_layer = maybeCurrentDynamicLayer();
  18901. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18902. int64_t cur_level = maybe_layer->layerId();
  18903. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  18904. return at::_ops::_foreach_clamp_min_List::call(self, other);
  18905. }
  18906. auto results = batch_rule(self, other);
  18907. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  18908. }
  18909. template <typename batch_rule_t, batch_rule_t batch_rule>
  18910. void _foreach_clamp_min__List_generated_plumbing(at::TensorList self, at::TensorList other) {
  18911. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18912. auto maybe_layer = maybeCurrentDynamicLayer();
  18913. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  18914. int64_t cur_level = maybe_layer->layerId();
  18915. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  18916. return at::_ops::_foreach_clamp_min__List::call(self, other);
  18917. }
  18918. batch_rule(self, other);
  18919. }
  18920. template <typename batch_rule_t, batch_rule_t batch_rule>
  18921. ::std::vector<at::Tensor> _foreach_clamp_min_ScalarList_generated_plumbing(at::TensorList self, at::ArrayRef<at::Scalar> scalars) {
  18922. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18923. auto maybe_layer = maybeCurrentDynamicLayer();
  18924. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18925. int64_t cur_level = maybe_layer->layerId();
  18926. if (!isBatchedAtLevel(self, cur_level)) {
  18927. return at::_ops::_foreach_clamp_min_ScalarList::call(self, scalars);
  18928. }
  18929. auto results = batch_rule(self, scalars);
  18930. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  18931. }
  18932. template <typename batch_rule_t, batch_rule_t batch_rule>
  18933. void _foreach_clamp_min__ScalarList_generated_plumbing(at::TensorList self, at::ArrayRef<at::Scalar> scalars) {
  18934. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18935. auto maybe_layer = maybeCurrentDynamicLayer();
  18936. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  18937. int64_t cur_level = maybe_layer->layerId();
  18938. if (!isBatchedAtLevel(self, cur_level)) {
  18939. return at::_ops::_foreach_clamp_min__ScalarList::call(self, scalars);
  18940. }
  18941. batch_rule(self, scalars);
  18942. }
  18943. template <typename batch_rule_t, batch_rule_t batch_rule>
  18944. ::std::vector<at::Tensor> _foreach_maximum_Scalar_generated_plumbing(at::TensorList self, const at::Scalar & scalar) {
  18945. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18946. auto maybe_layer = maybeCurrentDynamicLayer();
  18947. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18948. int64_t cur_level = maybe_layer->layerId();
  18949. if (!isBatchedAtLevel(self, cur_level)) {
  18950. return at::_ops::_foreach_maximum_Scalar::call(self, scalar);
  18951. }
  18952. auto results = batch_rule(self, scalar);
  18953. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  18954. }
  18955. template <typename batch_rule_t, batch_rule_t batch_rule>
  18956. void _foreach_maximum__Scalar_generated_plumbing(at::TensorList self, const at::Scalar & scalar) {
  18957. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18958. auto maybe_layer = maybeCurrentDynamicLayer();
  18959. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  18960. int64_t cur_level = maybe_layer->layerId();
  18961. if (!isBatchedAtLevel(self, cur_level)) {
  18962. return at::_ops::_foreach_maximum__Scalar::call(self, scalar);
  18963. }
  18964. batch_rule(self, scalar);
  18965. }
  18966. template <typename batch_rule_t, batch_rule_t batch_rule>
  18967. ::std::vector<at::Tensor> _foreach_maximum_List_generated_plumbing(at::TensorList self, at::TensorList other) {
  18968. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18969. auto maybe_layer = maybeCurrentDynamicLayer();
  18970. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18971. int64_t cur_level = maybe_layer->layerId();
  18972. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  18973. return at::_ops::_foreach_maximum_List::call(self, other);
  18974. }
  18975. auto results = batch_rule(self, other);
  18976. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  18977. }
  18978. template <typename batch_rule_t, batch_rule_t batch_rule>
  18979. void _foreach_maximum__List_generated_plumbing(at::TensorList self, at::TensorList other) {
  18980. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18981. auto maybe_layer = maybeCurrentDynamicLayer();
  18982. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  18983. int64_t cur_level = maybe_layer->layerId();
  18984. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  18985. return at::_ops::_foreach_maximum__List::call(self, other);
  18986. }
  18987. batch_rule(self, other);
  18988. }
  18989. template <typename batch_rule_t, batch_rule_t batch_rule>
  18990. ::std::vector<at::Tensor> _foreach_maximum_ScalarList_generated_plumbing(at::TensorList self, at::ArrayRef<at::Scalar> scalars) {
  18991. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  18992. auto maybe_layer = maybeCurrentDynamicLayer();
  18993. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  18994. int64_t cur_level = maybe_layer->layerId();
  18995. if (!isBatchedAtLevel(self, cur_level)) {
  18996. return at::_ops::_foreach_maximum_ScalarList::call(self, scalars);
  18997. }
  18998. auto results = batch_rule(self, scalars);
  18999. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19000. }
  19001. template <typename batch_rule_t, batch_rule_t batch_rule>
  19002. void _foreach_maximum__ScalarList_generated_plumbing(at::TensorList self, at::ArrayRef<at::Scalar> scalars) {
  19003. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19004. auto maybe_layer = maybeCurrentDynamicLayer();
  19005. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  19006. int64_t cur_level = maybe_layer->layerId();
  19007. if (!isBatchedAtLevel(self, cur_level)) {
  19008. return at::_ops::_foreach_maximum__ScalarList::call(self, scalars);
  19009. }
  19010. batch_rule(self, scalars);
  19011. }
  19012. template <typename batch_rule_t, batch_rule_t batch_rule>
  19013. ::std::vector<at::Tensor> _foreach_minimum_Scalar_generated_plumbing(at::TensorList self, const at::Scalar & scalar) {
  19014. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19015. auto maybe_layer = maybeCurrentDynamicLayer();
  19016. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19017. int64_t cur_level = maybe_layer->layerId();
  19018. if (!isBatchedAtLevel(self, cur_level)) {
  19019. return at::_ops::_foreach_minimum_Scalar::call(self, scalar);
  19020. }
  19021. auto results = batch_rule(self, scalar);
  19022. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19023. }
  19024. template <typename batch_rule_t, batch_rule_t batch_rule>
  19025. void _foreach_minimum__Scalar_generated_plumbing(at::TensorList self, const at::Scalar & scalar) {
  19026. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19027. auto maybe_layer = maybeCurrentDynamicLayer();
  19028. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  19029. int64_t cur_level = maybe_layer->layerId();
  19030. if (!isBatchedAtLevel(self, cur_level)) {
  19031. return at::_ops::_foreach_minimum__Scalar::call(self, scalar);
  19032. }
  19033. batch_rule(self, scalar);
  19034. }
  19035. template <typename batch_rule_t, batch_rule_t batch_rule>
  19036. ::std::vector<at::Tensor> _foreach_minimum_List_generated_plumbing(at::TensorList self, at::TensorList other) {
  19037. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19038. auto maybe_layer = maybeCurrentDynamicLayer();
  19039. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19040. int64_t cur_level = maybe_layer->layerId();
  19041. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  19042. return at::_ops::_foreach_minimum_List::call(self, other);
  19043. }
  19044. auto results = batch_rule(self, other);
  19045. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19046. }
  19047. template <typename batch_rule_t, batch_rule_t batch_rule>
  19048. void _foreach_minimum__List_generated_plumbing(at::TensorList self, at::TensorList other) {
  19049. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19050. auto maybe_layer = maybeCurrentDynamicLayer();
  19051. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  19052. int64_t cur_level = maybe_layer->layerId();
  19053. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  19054. return at::_ops::_foreach_minimum__List::call(self, other);
  19055. }
  19056. batch_rule(self, other);
  19057. }
  19058. template <typename batch_rule_t, batch_rule_t batch_rule>
  19059. ::std::vector<at::Tensor> _foreach_minimum_ScalarList_generated_plumbing(at::TensorList self, at::ArrayRef<at::Scalar> scalars) {
  19060. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19061. auto maybe_layer = maybeCurrentDynamicLayer();
  19062. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19063. int64_t cur_level = maybe_layer->layerId();
  19064. if (!isBatchedAtLevel(self, cur_level)) {
  19065. return at::_ops::_foreach_minimum_ScalarList::call(self, scalars);
  19066. }
  19067. auto results = batch_rule(self, scalars);
  19068. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19069. }
  19070. template <typename batch_rule_t, batch_rule_t batch_rule>
  19071. void _foreach_minimum__ScalarList_generated_plumbing(at::TensorList self, at::ArrayRef<at::Scalar> scalars) {
  19072. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19073. auto maybe_layer = maybeCurrentDynamicLayer();
  19074. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  19075. int64_t cur_level = maybe_layer->layerId();
  19076. if (!isBatchedAtLevel(self, cur_level)) {
  19077. return at::_ops::_foreach_minimum__ScalarList::call(self, scalars);
  19078. }
  19079. batch_rule(self, scalars);
  19080. }
  19081. template <typename batch_rule_t, batch_rule_t batch_rule>
  19082. ::std::vector<at::Tensor> _foreach_addcdiv_Scalar_generated_plumbing(at::TensorList self, at::TensorList tensor1, at::TensorList tensor2, const at::Scalar & value) {
  19083. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19084. auto maybe_layer = maybeCurrentDynamicLayer();
  19085. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19086. int64_t cur_level = maybe_layer->layerId();
  19087. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(tensor1, cur_level) && !isBatchedAtLevel(tensor2, cur_level)) {
  19088. return at::_ops::_foreach_addcdiv_Scalar::call(self, tensor1, tensor2, value);
  19089. }
  19090. auto results = batch_rule(self, tensor1, tensor2, value);
  19091. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19092. }
  19093. template <typename batch_rule_t, batch_rule_t batch_rule>
  19094. ::std::vector<at::Tensor> _foreach_addcdiv_ScalarList_generated_plumbing(at::TensorList self, at::TensorList tensor1, at::TensorList tensor2, at::ArrayRef<at::Scalar> scalars) {
  19095. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19096. auto maybe_layer = maybeCurrentDynamicLayer();
  19097. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19098. int64_t cur_level = maybe_layer->layerId();
  19099. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(tensor1, cur_level) && !isBatchedAtLevel(tensor2, cur_level)) {
  19100. return at::_ops::_foreach_addcdiv_ScalarList::call(self, tensor1, tensor2, scalars);
  19101. }
  19102. auto results = batch_rule(self, tensor1, tensor2, scalars);
  19103. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19104. }
  19105. template <typename batch_rule_t, batch_rule_t batch_rule>
  19106. ::std::vector<at::Tensor> _foreach_addcdiv_Tensor_generated_plumbing(at::TensorList self, at::TensorList tensor1, at::TensorList tensor2, const at::Tensor & scalars) {
  19107. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19108. auto maybe_layer = maybeCurrentDynamicLayer();
  19109. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19110. int64_t cur_level = maybe_layer->layerId();
  19111. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(tensor1, cur_level) && !isBatchedAtLevel(tensor2, cur_level) && !isBatchedAtLevel(scalars, cur_level)) {
  19112. return at::_ops::_foreach_addcdiv_Tensor::call(self, tensor1, tensor2, scalars);
  19113. }
  19114. auto [scalars_value, scalars_bdim] = unwrapTensorAtLevel(scalars, cur_level);
  19115. auto results = batch_rule(self, tensor1, tensor2, scalars_value, scalars_bdim);
  19116. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19117. }
  19118. template <typename batch_rule_t, batch_rule_t batch_rule>
  19119. void _foreach_addcdiv__Scalar_generated_plumbing(at::TensorList self, at::TensorList tensor1, at::TensorList tensor2, const at::Scalar & value) {
  19120. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19121. auto maybe_layer = maybeCurrentDynamicLayer();
  19122. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  19123. int64_t cur_level = maybe_layer->layerId();
  19124. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(tensor1, cur_level) && !isBatchedAtLevel(tensor2, cur_level)) {
  19125. return at::_ops::_foreach_addcdiv__Scalar::call(self, tensor1, tensor2, value);
  19126. }
  19127. batch_rule(self, tensor1, tensor2, value);
  19128. }
  19129. template <typename batch_rule_t, batch_rule_t batch_rule>
  19130. void _foreach_addcdiv__ScalarList_generated_plumbing(at::TensorList self, at::TensorList tensor1, at::TensorList tensor2, at::ArrayRef<at::Scalar> scalars) {
  19131. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19132. auto maybe_layer = maybeCurrentDynamicLayer();
  19133. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  19134. int64_t cur_level = maybe_layer->layerId();
  19135. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(tensor1, cur_level) && !isBatchedAtLevel(tensor2, cur_level)) {
  19136. return at::_ops::_foreach_addcdiv__ScalarList::call(self, tensor1, tensor2, scalars);
  19137. }
  19138. batch_rule(self, tensor1, tensor2, scalars);
  19139. }
  19140. template <typename batch_rule_t, batch_rule_t batch_rule>
  19141. void _foreach_addcdiv__Tensor_generated_plumbing(at::TensorList self, at::TensorList tensor1, at::TensorList tensor2, const at::Tensor & scalars) {
  19142. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19143. auto maybe_layer = maybeCurrentDynamicLayer();
  19144. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  19145. int64_t cur_level = maybe_layer->layerId();
  19146. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(tensor1, cur_level) && !isBatchedAtLevel(tensor2, cur_level) && !isBatchedAtLevel(scalars, cur_level)) {
  19147. return at::_ops::_foreach_addcdiv__Tensor::call(self, tensor1, tensor2, scalars);
  19148. }
  19149. auto [scalars_value, scalars_bdim] = unwrapTensorAtLevel(scalars, cur_level);
  19150. batch_rule(self, tensor1, tensor2, scalars_value, scalars_bdim);
  19151. }
  19152. template <typename batch_rule_t, batch_rule_t batch_rule>
  19153. ::std::vector<at::Tensor> _foreach_addcmul_Scalar_generated_plumbing(at::TensorList self, at::TensorList tensor1, at::TensorList tensor2, const at::Scalar & value) {
  19154. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19155. auto maybe_layer = maybeCurrentDynamicLayer();
  19156. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19157. int64_t cur_level = maybe_layer->layerId();
  19158. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(tensor1, cur_level) && !isBatchedAtLevel(tensor2, cur_level)) {
  19159. return at::_ops::_foreach_addcmul_Scalar::call(self, tensor1, tensor2, value);
  19160. }
  19161. auto results = batch_rule(self, tensor1, tensor2, value);
  19162. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19163. }
  19164. template <typename batch_rule_t, batch_rule_t batch_rule>
  19165. ::std::vector<at::Tensor> _foreach_addcmul_ScalarList_generated_plumbing(at::TensorList self, at::TensorList tensor1, at::TensorList tensor2, at::ArrayRef<at::Scalar> scalars) {
  19166. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19167. auto maybe_layer = maybeCurrentDynamicLayer();
  19168. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19169. int64_t cur_level = maybe_layer->layerId();
  19170. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(tensor1, cur_level) && !isBatchedAtLevel(tensor2, cur_level)) {
  19171. return at::_ops::_foreach_addcmul_ScalarList::call(self, tensor1, tensor2, scalars);
  19172. }
  19173. auto results = batch_rule(self, tensor1, tensor2, scalars);
  19174. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19175. }
  19176. template <typename batch_rule_t, batch_rule_t batch_rule>
  19177. ::std::vector<at::Tensor> _foreach_addcmul_Tensor_generated_plumbing(at::TensorList self, at::TensorList tensor1, at::TensorList tensor2, const at::Tensor & scalars) {
  19178. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19179. auto maybe_layer = maybeCurrentDynamicLayer();
  19180. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19181. int64_t cur_level = maybe_layer->layerId();
  19182. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(tensor1, cur_level) && !isBatchedAtLevel(tensor2, cur_level) && !isBatchedAtLevel(scalars, cur_level)) {
  19183. return at::_ops::_foreach_addcmul_Tensor::call(self, tensor1, tensor2, scalars);
  19184. }
  19185. auto [scalars_value, scalars_bdim] = unwrapTensorAtLevel(scalars, cur_level);
  19186. auto results = batch_rule(self, tensor1, tensor2, scalars_value, scalars_bdim);
  19187. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19188. }
  19189. template <typename batch_rule_t, batch_rule_t batch_rule>
  19190. void _foreach_addcmul__Scalar_generated_plumbing(at::TensorList self, at::TensorList tensor1, at::TensorList tensor2, const at::Scalar & value) {
  19191. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19192. auto maybe_layer = maybeCurrentDynamicLayer();
  19193. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  19194. int64_t cur_level = maybe_layer->layerId();
  19195. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(tensor1, cur_level) && !isBatchedAtLevel(tensor2, cur_level)) {
  19196. return at::_ops::_foreach_addcmul__Scalar::call(self, tensor1, tensor2, value);
  19197. }
  19198. batch_rule(self, tensor1, tensor2, value);
  19199. }
  19200. template <typename batch_rule_t, batch_rule_t batch_rule>
  19201. void _foreach_addcmul__ScalarList_generated_plumbing(at::TensorList self, at::TensorList tensor1, at::TensorList tensor2, at::ArrayRef<at::Scalar> scalars) {
  19202. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19203. auto maybe_layer = maybeCurrentDynamicLayer();
  19204. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  19205. int64_t cur_level = maybe_layer->layerId();
  19206. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(tensor1, cur_level) && !isBatchedAtLevel(tensor2, cur_level)) {
  19207. return at::_ops::_foreach_addcmul__ScalarList::call(self, tensor1, tensor2, scalars);
  19208. }
  19209. batch_rule(self, tensor1, tensor2, scalars);
  19210. }
  19211. template <typename batch_rule_t, batch_rule_t batch_rule>
  19212. void _foreach_addcmul__Tensor_generated_plumbing(at::TensorList self, at::TensorList tensor1, at::TensorList tensor2, const at::Tensor & scalars) {
  19213. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19214. auto maybe_layer = maybeCurrentDynamicLayer();
  19215. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  19216. int64_t cur_level = maybe_layer->layerId();
  19217. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(tensor1, cur_level) && !isBatchedAtLevel(tensor2, cur_level) && !isBatchedAtLevel(scalars, cur_level)) {
  19218. return at::_ops::_foreach_addcmul__Tensor::call(self, tensor1, tensor2, scalars);
  19219. }
  19220. auto [scalars_value, scalars_bdim] = unwrapTensorAtLevel(scalars, cur_level);
  19221. batch_rule(self, tensor1, tensor2, scalars_value, scalars_bdim);
  19222. }
  19223. template <typename batch_rule_t, batch_rule_t batch_rule>
  19224. ::std::vector<at::Tensor> _foreach_abs_generated_plumbing(at::TensorList self) {
  19225. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19226. auto maybe_layer = maybeCurrentDynamicLayer();
  19227. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19228. int64_t cur_level = maybe_layer->layerId();
  19229. if (!isBatchedAtLevel(self, cur_level)) {
  19230. return at::_ops::_foreach_abs::call(self);
  19231. }
  19232. auto results = batch_rule(self);
  19233. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19234. }
  19235. template <typename batch_rule_t, batch_rule_t batch_rule>
  19236. void _foreach_abs__generated_plumbing(at::TensorList self) {
  19237. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19238. auto maybe_layer = maybeCurrentDynamicLayer();
  19239. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  19240. int64_t cur_level = maybe_layer->layerId();
  19241. if (!isBatchedAtLevel(self, cur_level)) {
  19242. return at::_ops::_foreach_abs_::call(self);
  19243. }
  19244. batch_rule(self);
  19245. }
  19246. template <typename batch_rule_t, batch_rule_t batch_rule>
  19247. ::std::vector<at::Tensor> _foreach_acos_generated_plumbing(at::TensorList self) {
  19248. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19249. auto maybe_layer = maybeCurrentDynamicLayer();
  19250. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19251. int64_t cur_level = maybe_layer->layerId();
  19252. if (!isBatchedAtLevel(self, cur_level)) {
  19253. return at::_ops::_foreach_acos::call(self);
  19254. }
  19255. auto results = batch_rule(self);
  19256. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19257. }
  19258. template <typename batch_rule_t, batch_rule_t batch_rule>
  19259. void _foreach_acos__generated_plumbing(at::TensorList self) {
  19260. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19261. auto maybe_layer = maybeCurrentDynamicLayer();
  19262. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  19263. int64_t cur_level = maybe_layer->layerId();
  19264. if (!isBatchedAtLevel(self, cur_level)) {
  19265. return at::_ops::_foreach_acos_::call(self);
  19266. }
  19267. batch_rule(self);
  19268. }
  19269. template <typename batch_rule_t, batch_rule_t batch_rule>
  19270. ::std::vector<at::Tensor> _foreach_asin_generated_plumbing(at::TensorList self) {
  19271. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19272. auto maybe_layer = maybeCurrentDynamicLayer();
  19273. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19274. int64_t cur_level = maybe_layer->layerId();
  19275. if (!isBatchedAtLevel(self, cur_level)) {
  19276. return at::_ops::_foreach_asin::call(self);
  19277. }
  19278. auto results = batch_rule(self);
  19279. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19280. }
  19281. template <typename batch_rule_t, batch_rule_t batch_rule>
  19282. void _foreach_asin__generated_plumbing(at::TensorList self) {
  19283. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19284. auto maybe_layer = maybeCurrentDynamicLayer();
  19285. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  19286. int64_t cur_level = maybe_layer->layerId();
  19287. if (!isBatchedAtLevel(self, cur_level)) {
  19288. return at::_ops::_foreach_asin_::call(self);
  19289. }
  19290. batch_rule(self);
  19291. }
  19292. template <typename batch_rule_t, batch_rule_t batch_rule>
  19293. ::std::vector<at::Tensor> _foreach_atan_generated_plumbing(at::TensorList self) {
  19294. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19295. auto maybe_layer = maybeCurrentDynamicLayer();
  19296. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19297. int64_t cur_level = maybe_layer->layerId();
  19298. if (!isBatchedAtLevel(self, cur_level)) {
  19299. return at::_ops::_foreach_atan::call(self);
  19300. }
  19301. auto results = batch_rule(self);
  19302. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19303. }
  19304. template <typename batch_rule_t, batch_rule_t batch_rule>
  19305. void _foreach_atan__generated_plumbing(at::TensorList self) {
  19306. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19307. auto maybe_layer = maybeCurrentDynamicLayer();
  19308. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  19309. int64_t cur_level = maybe_layer->layerId();
  19310. if (!isBatchedAtLevel(self, cur_level)) {
  19311. return at::_ops::_foreach_atan_::call(self);
  19312. }
  19313. batch_rule(self);
  19314. }
  19315. template <typename batch_rule_t, batch_rule_t batch_rule>
  19316. ::std::vector<at::Tensor> _foreach_ceil_generated_plumbing(at::TensorList self) {
  19317. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19318. auto maybe_layer = maybeCurrentDynamicLayer();
  19319. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19320. int64_t cur_level = maybe_layer->layerId();
  19321. if (!isBatchedAtLevel(self, cur_level)) {
  19322. return at::_ops::_foreach_ceil::call(self);
  19323. }
  19324. auto results = batch_rule(self);
  19325. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19326. }
  19327. template <typename batch_rule_t, batch_rule_t batch_rule>
  19328. void _foreach_ceil__generated_plumbing(at::TensorList self) {
  19329. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19330. auto maybe_layer = maybeCurrentDynamicLayer();
  19331. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  19332. int64_t cur_level = maybe_layer->layerId();
  19333. if (!isBatchedAtLevel(self, cur_level)) {
  19334. return at::_ops::_foreach_ceil_::call(self);
  19335. }
  19336. batch_rule(self);
  19337. }
  19338. template <typename batch_rule_t, batch_rule_t batch_rule>
  19339. ::std::vector<at::Tensor> _foreach_cos_generated_plumbing(at::TensorList self) {
  19340. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19341. auto maybe_layer = maybeCurrentDynamicLayer();
  19342. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19343. int64_t cur_level = maybe_layer->layerId();
  19344. if (!isBatchedAtLevel(self, cur_level)) {
  19345. return at::_ops::_foreach_cos::call(self);
  19346. }
  19347. auto results = batch_rule(self);
  19348. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19349. }
  19350. template <typename batch_rule_t, batch_rule_t batch_rule>
  19351. void _foreach_cos__generated_plumbing(at::TensorList self) {
  19352. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19353. auto maybe_layer = maybeCurrentDynamicLayer();
  19354. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  19355. int64_t cur_level = maybe_layer->layerId();
  19356. if (!isBatchedAtLevel(self, cur_level)) {
  19357. return at::_ops::_foreach_cos_::call(self);
  19358. }
  19359. batch_rule(self);
  19360. }
  19361. template <typename batch_rule_t, batch_rule_t batch_rule>
  19362. ::std::vector<at::Tensor> _foreach_cosh_generated_plumbing(at::TensorList self) {
  19363. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19364. auto maybe_layer = maybeCurrentDynamicLayer();
  19365. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19366. int64_t cur_level = maybe_layer->layerId();
  19367. if (!isBatchedAtLevel(self, cur_level)) {
  19368. return at::_ops::_foreach_cosh::call(self);
  19369. }
  19370. auto results = batch_rule(self);
  19371. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19372. }
  19373. template <typename batch_rule_t, batch_rule_t batch_rule>
  19374. void _foreach_cosh__generated_plumbing(at::TensorList self) {
  19375. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19376. auto maybe_layer = maybeCurrentDynamicLayer();
  19377. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  19378. int64_t cur_level = maybe_layer->layerId();
  19379. if (!isBatchedAtLevel(self, cur_level)) {
  19380. return at::_ops::_foreach_cosh_::call(self);
  19381. }
  19382. batch_rule(self);
  19383. }
  19384. template <typename batch_rule_t, batch_rule_t batch_rule>
  19385. ::std::vector<at::Tensor> _foreach_erf_generated_plumbing(at::TensorList self) {
  19386. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19387. auto maybe_layer = maybeCurrentDynamicLayer();
  19388. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19389. int64_t cur_level = maybe_layer->layerId();
  19390. if (!isBatchedAtLevel(self, cur_level)) {
  19391. return at::_ops::_foreach_erf::call(self);
  19392. }
  19393. auto results = batch_rule(self);
  19394. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19395. }
  19396. template <typename batch_rule_t, batch_rule_t batch_rule>
  19397. void _foreach_erf__generated_plumbing(at::TensorList self) {
  19398. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19399. auto maybe_layer = maybeCurrentDynamicLayer();
  19400. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  19401. int64_t cur_level = maybe_layer->layerId();
  19402. if (!isBatchedAtLevel(self, cur_level)) {
  19403. return at::_ops::_foreach_erf_::call(self);
  19404. }
  19405. batch_rule(self);
  19406. }
  19407. template <typename batch_rule_t, batch_rule_t batch_rule>
  19408. ::std::vector<at::Tensor> _foreach_erfc_generated_plumbing(at::TensorList self) {
  19409. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19410. auto maybe_layer = maybeCurrentDynamicLayer();
  19411. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19412. int64_t cur_level = maybe_layer->layerId();
  19413. if (!isBatchedAtLevel(self, cur_level)) {
  19414. return at::_ops::_foreach_erfc::call(self);
  19415. }
  19416. auto results = batch_rule(self);
  19417. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19418. }
  19419. template <typename batch_rule_t, batch_rule_t batch_rule>
  19420. void _foreach_erfc__generated_plumbing(at::TensorList self) {
  19421. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19422. auto maybe_layer = maybeCurrentDynamicLayer();
  19423. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  19424. int64_t cur_level = maybe_layer->layerId();
  19425. if (!isBatchedAtLevel(self, cur_level)) {
  19426. return at::_ops::_foreach_erfc_::call(self);
  19427. }
  19428. batch_rule(self);
  19429. }
  19430. template <typename batch_rule_t, batch_rule_t batch_rule>
  19431. ::std::vector<at::Tensor> _foreach_exp_generated_plumbing(at::TensorList self) {
  19432. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19433. auto maybe_layer = maybeCurrentDynamicLayer();
  19434. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19435. int64_t cur_level = maybe_layer->layerId();
  19436. if (!isBatchedAtLevel(self, cur_level)) {
  19437. return at::_ops::_foreach_exp::call(self);
  19438. }
  19439. auto results = batch_rule(self);
  19440. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19441. }
  19442. template <typename batch_rule_t, batch_rule_t batch_rule>
  19443. void _foreach_exp__generated_plumbing(at::TensorList self) {
  19444. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19445. auto maybe_layer = maybeCurrentDynamicLayer();
  19446. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  19447. int64_t cur_level = maybe_layer->layerId();
  19448. if (!isBatchedAtLevel(self, cur_level)) {
  19449. return at::_ops::_foreach_exp_::call(self);
  19450. }
  19451. batch_rule(self);
  19452. }
  19453. template <typename batch_rule_t, batch_rule_t batch_rule>
  19454. ::std::vector<at::Tensor> _foreach_expm1_generated_plumbing(at::TensorList self) {
  19455. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19456. auto maybe_layer = maybeCurrentDynamicLayer();
  19457. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19458. int64_t cur_level = maybe_layer->layerId();
  19459. if (!isBatchedAtLevel(self, cur_level)) {
  19460. return at::_ops::_foreach_expm1::call(self);
  19461. }
  19462. auto results = batch_rule(self);
  19463. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19464. }
  19465. template <typename batch_rule_t, batch_rule_t batch_rule>
  19466. void _foreach_expm1__generated_plumbing(at::TensorList self) {
  19467. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19468. auto maybe_layer = maybeCurrentDynamicLayer();
  19469. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  19470. int64_t cur_level = maybe_layer->layerId();
  19471. if (!isBatchedAtLevel(self, cur_level)) {
  19472. return at::_ops::_foreach_expm1_::call(self);
  19473. }
  19474. batch_rule(self);
  19475. }
  19476. template <typename batch_rule_t, batch_rule_t batch_rule>
  19477. ::std::vector<at::Tensor> _foreach_floor_generated_plumbing(at::TensorList self) {
  19478. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19479. auto maybe_layer = maybeCurrentDynamicLayer();
  19480. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19481. int64_t cur_level = maybe_layer->layerId();
  19482. if (!isBatchedAtLevel(self, cur_level)) {
  19483. return at::_ops::_foreach_floor::call(self);
  19484. }
  19485. auto results = batch_rule(self);
  19486. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19487. }
  19488. template <typename batch_rule_t, batch_rule_t batch_rule>
  19489. void _foreach_floor__generated_plumbing(at::TensorList self) {
  19490. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19491. auto maybe_layer = maybeCurrentDynamicLayer();
  19492. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  19493. int64_t cur_level = maybe_layer->layerId();
  19494. if (!isBatchedAtLevel(self, cur_level)) {
  19495. return at::_ops::_foreach_floor_::call(self);
  19496. }
  19497. batch_rule(self);
  19498. }
  19499. template <typename batch_rule_t, batch_rule_t batch_rule>
  19500. ::std::vector<at::Tensor> _foreach_frac_generated_plumbing(at::TensorList self) {
  19501. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19502. auto maybe_layer = maybeCurrentDynamicLayer();
  19503. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19504. int64_t cur_level = maybe_layer->layerId();
  19505. if (!isBatchedAtLevel(self, cur_level)) {
  19506. return at::_ops::_foreach_frac::call(self);
  19507. }
  19508. auto results = batch_rule(self);
  19509. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19510. }
  19511. template <typename batch_rule_t, batch_rule_t batch_rule>
  19512. void _foreach_frac__generated_plumbing(at::TensorList self) {
  19513. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19514. auto maybe_layer = maybeCurrentDynamicLayer();
  19515. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  19516. int64_t cur_level = maybe_layer->layerId();
  19517. if (!isBatchedAtLevel(self, cur_level)) {
  19518. return at::_ops::_foreach_frac_::call(self);
  19519. }
  19520. batch_rule(self);
  19521. }
  19522. template <typename batch_rule_t, batch_rule_t batch_rule>
  19523. ::std::vector<at::Tensor> _foreach_lerp_List_generated_plumbing(at::TensorList self, at::TensorList tensors1, at::TensorList weights) {
  19524. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19525. auto maybe_layer = maybeCurrentDynamicLayer();
  19526. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19527. int64_t cur_level = maybe_layer->layerId();
  19528. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(tensors1, cur_level) && !isBatchedAtLevel(weights, cur_level)) {
  19529. return at::_ops::_foreach_lerp_List::call(self, tensors1, weights);
  19530. }
  19531. auto results = batch_rule(self, tensors1, weights);
  19532. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19533. }
  19534. template <typename batch_rule_t, batch_rule_t batch_rule>
  19535. void _foreach_lerp__List_generated_plumbing(at::TensorList self, at::TensorList tensors1, at::TensorList weights) {
  19536. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19537. auto maybe_layer = maybeCurrentDynamicLayer();
  19538. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  19539. int64_t cur_level = maybe_layer->layerId();
  19540. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(tensors1, cur_level) && !isBatchedAtLevel(weights, cur_level)) {
  19541. return at::_ops::_foreach_lerp__List::call(self, tensors1, weights);
  19542. }
  19543. batch_rule(self, tensors1, weights);
  19544. }
  19545. template <typename batch_rule_t, batch_rule_t batch_rule>
  19546. ::std::vector<at::Tensor> _foreach_lerp_Scalar_generated_plumbing(at::TensorList self, at::TensorList tensors1, const at::Scalar & weight) {
  19547. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19548. auto maybe_layer = maybeCurrentDynamicLayer();
  19549. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19550. int64_t cur_level = maybe_layer->layerId();
  19551. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(tensors1, cur_level)) {
  19552. return at::_ops::_foreach_lerp_Scalar::call(self, tensors1, weight);
  19553. }
  19554. auto results = batch_rule(self, tensors1, weight);
  19555. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19556. }
  19557. template <typename batch_rule_t, batch_rule_t batch_rule>
  19558. void _foreach_lerp__Scalar_generated_plumbing(at::TensorList self, at::TensorList tensors1, const at::Scalar & weight) {
  19559. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19560. auto maybe_layer = maybeCurrentDynamicLayer();
  19561. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  19562. int64_t cur_level = maybe_layer->layerId();
  19563. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(tensors1, cur_level)) {
  19564. return at::_ops::_foreach_lerp__Scalar::call(self, tensors1, weight);
  19565. }
  19566. batch_rule(self, tensors1, weight);
  19567. }
  19568. template <typename batch_rule_t, batch_rule_t batch_rule>
  19569. ::std::vector<at::Tensor> _foreach_lerp_ScalarList_generated_plumbing(at::TensorList self, at::TensorList tensors1, at::ArrayRef<at::Scalar> weight) {
  19570. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19571. auto maybe_layer = maybeCurrentDynamicLayer();
  19572. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19573. int64_t cur_level = maybe_layer->layerId();
  19574. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(tensors1, cur_level)) {
  19575. return at::_ops::_foreach_lerp_ScalarList::call(self, tensors1, weight);
  19576. }
  19577. auto results = batch_rule(self, tensors1, weight);
  19578. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19579. }
  19580. template <typename batch_rule_t, batch_rule_t batch_rule>
  19581. void _foreach_lerp__ScalarList_generated_plumbing(at::TensorList self, at::TensorList tensors1, at::ArrayRef<at::Scalar> weight) {
  19582. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19583. auto maybe_layer = maybeCurrentDynamicLayer();
  19584. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  19585. int64_t cur_level = maybe_layer->layerId();
  19586. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(tensors1, cur_level)) {
  19587. return at::_ops::_foreach_lerp__ScalarList::call(self, tensors1, weight);
  19588. }
  19589. batch_rule(self, tensors1, weight);
  19590. }
  19591. template <typename batch_rule_t, batch_rule_t batch_rule>
  19592. ::std::vector<at::Tensor> _foreach_lgamma_generated_plumbing(at::TensorList self) {
  19593. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19594. auto maybe_layer = maybeCurrentDynamicLayer();
  19595. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19596. int64_t cur_level = maybe_layer->layerId();
  19597. if (!isBatchedAtLevel(self, cur_level)) {
  19598. return at::_ops::_foreach_lgamma::call(self);
  19599. }
  19600. auto results = batch_rule(self);
  19601. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19602. }
  19603. template <typename batch_rule_t, batch_rule_t batch_rule>
  19604. void _foreach_lgamma__generated_plumbing(at::TensorList self) {
  19605. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19606. auto maybe_layer = maybeCurrentDynamicLayer();
  19607. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  19608. int64_t cur_level = maybe_layer->layerId();
  19609. if (!isBatchedAtLevel(self, cur_level)) {
  19610. return at::_ops::_foreach_lgamma_::call(self);
  19611. }
  19612. batch_rule(self);
  19613. }
  19614. template <typename batch_rule_t, batch_rule_t batch_rule>
  19615. ::std::vector<at::Tensor> _foreach_log_generated_plumbing(at::TensorList self) {
  19616. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19617. auto maybe_layer = maybeCurrentDynamicLayer();
  19618. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19619. int64_t cur_level = maybe_layer->layerId();
  19620. if (!isBatchedAtLevel(self, cur_level)) {
  19621. return at::_ops::_foreach_log::call(self);
  19622. }
  19623. auto results = batch_rule(self);
  19624. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19625. }
  19626. template <typename batch_rule_t, batch_rule_t batch_rule>
  19627. void _foreach_log__generated_plumbing(at::TensorList self) {
  19628. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19629. auto maybe_layer = maybeCurrentDynamicLayer();
  19630. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  19631. int64_t cur_level = maybe_layer->layerId();
  19632. if (!isBatchedAtLevel(self, cur_level)) {
  19633. return at::_ops::_foreach_log_::call(self);
  19634. }
  19635. batch_rule(self);
  19636. }
  19637. template <typename batch_rule_t, batch_rule_t batch_rule>
  19638. ::std::vector<at::Tensor> _foreach_log10_generated_plumbing(at::TensorList self) {
  19639. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19640. auto maybe_layer = maybeCurrentDynamicLayer();
  19641. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19642. int64_t cur_level = maybe_layer->layerId();
  19643. if (!isBatchedAtLevel(self, cur_level)) {
  19644. return at::_ops::_foreach_log10::call(self);
  19645. }
  19646. auto results = batch_rule(self);
  19647. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19648. }
  19649. template <typename batch_rule_t, batch_rule_t batch_rule>
  19650. void _foreach_log10__generated_plumbing(at::TensorList self) {
  19651. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19652. auto maybe_layer = maybeCurrentDynamicLayer();
  19653. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  19654. int64_t cur_level = maybe_layer->layerId();
  19655. if (!isBatchedAtLevel(self, cur_level)) {
  19656. return at::_ops::_foreach_log10_::call(self);
  19657. }
  19658. batch_rule(self);
  19659. }
  19660. template <typename batch_rule_t, batch_rule_t batch_rule>
  19661. ::std::vector<at::Tensor> _foreach_log1p_generated_plumbing(at::TensorList self) {
  19662. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19663. auto maybe_layer = maybeCurrentDynamicLayer();
  19664. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19665. int64_t cur_level = maybe_layer->layerId();
  19666. if (!isBatchedAtLevel(self, cur_level)) {
  19667. return at::_ops::_foreach_log1p::call(self);
  19668. }
  19669. auto results = batch_rule(self);
  19670. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19671. }
  19672. template <typename batch_rule_t, batch_rule_t batch_rule>
  19673. void _foreach_log1p__generated_plumbing(at::TensorList self) {
  19674. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19675. auto maybe_layer = maybeCurrentDynamicLayer();
  19676. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  19677. int64_t cur_level = maybe_layer->layerId();
  19678. if (!isBatchedAtLevel(self, cur_level)) {
  19679. return at::_ops::_foreach_log1p_::call(self);
  19680. }
  19681. batch_rule(self);
  19682. }
  19683. template <typename batch_rule_t, batch_rule_t batch_rule>
  19684. ::std::vector<at::Tensor> _foreach_log2_generated_plumbing(at::TensorList self) {
  19685. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19686. auto maybe_layer = maybeCurrentDynamicLayer();
  19687. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19688. int64_t cur_level = maybe_layer->layerId();
  19689. if (!isBatchedAtLevel(self, cur_level)) {
  19690. return at::_ops::_foreach_log2::call(self);
  19691. }
  19692. auto results = batch_rule(self);
  19693. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19694. }
  19695. template <typename batch_rule_t, batch_rule_t batch_rule>
  19696. void _foreach_log2__generated_plumbing(at::TensorList self) {
  19697. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19698. auto maybe_layer = maybeCurrentDynamicLayer();
  19699. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  19700. int64_t cur_level = maybe_layer->layerId();
  19701. if (!isBatchedAtLevel(self, cur_level)) {
  19702. return at::_ops::_foreach_log2_::call(self);
  19703. }
  19704. batch_rule(self);
  19705. }
  19706. template <typename batch_rule_t, batch_rule_t batch_rule>
  19707. ::std::vector<at::Tensor> _foreach_max_generated_plumbing(at::TensorList self) {
  19708. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19709. auto maybe_layer = maybeCurrentDynamicLayer();
  19710. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19711. int64_t cur_level = maybe_layer->layerId();
  19712. if (!isBatchedAtLevel(self, cur_level)) {
  19713. return at::_ops::_foreach_max::call(self);
  19714. }
  19715. auto results = batch_rule(self);
  19716. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19717. }
  19718. template <typename batch_rule_t, batch_rule_t batch_rule>
  19719. ::std::vector<at::Tensor> _foreach_neg_generated_plumbing(at::TensorList self) {
  19720. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19721. auto maybe_layer = maybeCurrentDynamicLayer();
  19722. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19723. int64_t cur_level = maybe_layer->layerId();
  19724. if (!isBatchedAtLevel(self, cur_level)) {
  19725. return at::_ops::_foreach_neg::call(self);
  19726. }
  19727. auto results = batch_rule(self);
  19728. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19729. }
  19730. template <typename batch_rule_t, batch_rule_t batch_rule>
  19731. void _foreach_neg__generated_plumbing(at::TensorList self) {
  19732. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19733. auto maybe_layer = maybeCurrentDynamicLayer();
  19734. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  19735. int64_t cur_level = maybe_layer->layerId();
  19736. if (!isBatchedAtLevel(self, cur_level)) {
  19737. return at::_ops::_foreach_neg_::call(self);
  19738. }
  19739. batch_rule(self);
  19740. }
  19741. template <typename batch_rule_t, batch_rule_t batch_rule>
  19742. ::std::vector<at::Tensor> _foreach_norm_Scalar_generated_plumbing(at::TensorList self, const at::Scalar & ord, ::std::optional<at::ScalarType> dtype) {
  19743. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19744. auto maybe_layer = maybeCurrentDynamicLayer();
  19745. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19746. int64_t cur_level = maybe_layer->layerId();
  19747. if (!isBatchedAtLevel(self, cur_level)) {
  19748. return at::_ops::_foreach_norm_Scalar::call(self, ord, dtype);
  19749. }
  19750. auto results = batch_rule(self, ord, dtype);
  19751. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19752. }
  19753. template <typename batch_rule_t, batch_rule_t batch_rule>
  19754. ::std::vector<at::Tensor> _foreach_pow_List_generated_plumbing(at::TensorList self, at::TensorList exponent) {
  19755. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19756. auto maybe_layer = maybeCurrentDynamicLayer();
  19757. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19758. int64_t cur_level = maybe_layer->layerId();
  19759. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(exponent, cur_level)) {
  19760. return at::_ops::_foreach_pow_List::call(self, exponent);
  19761. }
  19762. auto results = batch_rule(self, exponent);
  19763. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19764. }
  19765. template <typename batch_rule_t, batch_rule_t batch_rule>
  19766. ::std::vector<at::Tensor> _foreach_pow_Scalar_generated_plumbing(at::TensorList self, const at::Scalar & exponent) {
  19767. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19768. auto maybe_layer = maybeCurrentDynamicLayer();
  19769. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19770. int64_t cur_level = maybe_layer->layerId();
  19771. if (!isBatchedAtLevel(self, cur_level)) {
  19772. return at::_ops::_foreach_pow_Scalar::call(self, exponent);
  19773. }
  19774. auto results = batch_rule(self, exponent);
  19775. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19776. }
  19777. template <typename batch_rule_t, batch_rule_t batch_rule>
  19778. ::std::vector<at::Tensor> _foreach_pow_ScalarList_generated_plumbing(at::TensorList self, at::ArrayRef<at::Scalar> exponent) {
  19779. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19780. auto maybe_layer = maybeCurrentDynamicLayer();
  19781. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19782. int64_t cur_level = maybe_layer->layerId();
  19783. if (!isBatchedAtLevel(self, cur_level)) {
  19784. return at::_ops::_foreach_pow_ScalarList::call(self, exponent);
  19785. }
  19786. auto results = batch_rule(self, exponent);
  19787. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19788. }
  19789. template <typename batch_rule_t, batch_rule_t batch_rule>
  19790. ::std::vector<at::Tensor> _foreach_pow_ScalarAndTensor_generated_plumbing(const at::Scalar & self, at::TensorList exponent) {
  19791. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19792. auto maybe_layer = maybeCurrentDynamicLayer();
  19793. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19794. int64_t cur_level = maybe_layer->layerId();
  19795. if (!isBatchedAtLevel(exponent, cur_level)) {
  19796. return at::_ops::_foreach_pow_ScalarAndTensor::call(self, exponent);
  19797. }
  19798. auto results = batch_rule(self, exponent);
  19799. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19800. }
  19801. template <typename batch_rule_t, batch_rule_t batch_rule>
  19802. void _foreach_pow__List_generated_plumbing(at::TensorList self, at::TensorList exponent) {
  19803. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19804. auto maybe_layer = maybeCurrentDynamicLayer();
  19805. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  19806. int64_t cur_level = maybe_layer->layerId();
  19807. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(exponent, cur_level)) {
  19808. return at::_ops::_foreach_pow__List::call(self, exponent);
  19809. }
  19810. batch_rule(self, exponent);
  19811. }
  19812. template <typename batch_rule_t, batch_rule_t batch_rule>
  19813. void _foreach_pow__Scalar_generated_plumbing(at::TensorList self, const at::Scalar & exponent) {
  19814. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19815. auto maybe_layer = maybeCurrentDynamicLayer();
  19816. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  19817. int64_t cur_level = maybe_layer->layerId();
  19818. if (!isBatchedAtLevel(self, cur_level)) {
  19819. return at::_ops::_foreach_pow__Scalar::call(self, exponent);
  19820. }
  19821. batch_rule(self, exponent);
  19822. }
  19823. template <typename batch_rule_t, batch_rule_t batch_rule>
  19824. void _foreach_pow__ScalarList_generated_plumbing(at::TensorList self, at::ArrayRef<at::Scalar> exponent) {
  19825. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19826. auto maybe_layer = maybeCurrentDynamicLayer();
  19827. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  19828. int64_t cur_level = maybe_layer->layerId();
  19829. if (!isBatchedAtLevel(self, cur_level)) {
  19830. return at::_ops::_foreach_pow__ScalarList::call(self, exponent);
  19831. }
  19832. batch_rule(self, exponent);
  19833. }
  19834. template <typename batch_rule_t, batch_rule_t batch_rule>
  19835. ::std::vector<at::Tensor> _foreach_reciprocal_generated_plumbing(at::TensorList self) {
  19836. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19837. auto maybe_layer = maybeCurrentDynamicLayer();
  19838. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19839. int64_t cur_level = maybe_layer->layerId();
  19840. if (!isBatchedAtLevel(self, cur_level)) {
  19841. return at::_ops::_foreach_reciprocal::call(self);
  19842. }
  19843. auto results = batch_rule(self);
  19844. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19845. }
  19846. template <typename batch_rule_t, batch_rule_t batch_rule>
  19847. void _foreach_reciprocal__generated_plumbing(at::TensorList self) {
  19848. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19849. auto maybe_layer = maybeCurrentDynamicLayer();
  19850. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  19851. int64_t cur_level = maybe_layer->layerId();
  19852. if (!isBatchedAtLevel(self, cur_level)) {
  19853. return at::_ops::_foreach_reciprocal_::call(self);
  19854. }
  19855. batch_rule(self);
  19856. }
  19857. template <typename batch_rule_t, batch_rule_t batch_rule>
  19858. ::std::vector<at::Tensor> _foreach_round_generated_plumbing(at::TensorList self) {
  19859. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19860. auto maybe_layer = maybeCurrentDynamicLayer();
  19861. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19862. int64_t cur_level = maybe_layer->layerId();
  19863. if (!isBatchedAtLevel(self, cur_level)) {
  19864. return at::_ops::_foreach_round::call(self);
  19865. }
  19866. auto results = batch_rule(self);
  19867. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19868. }
  19869. template <typename batch_rule_t, batch_rule_t batch_rule>
  19870. void _foreach_round__generated_plumbing(at::TensorList self) {
  19871. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19872. auto maybe_layer = maybeCurrentDynamicLayer();
  19873. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  19874. int64_t cur_level = maybe_layer->layerId();
  19875. if (!isBatchedAtLevel(self, cur_level)) {
  19876. return at::_ops::_foreach_round_::call(self);
  19877. }
  19878. batch_rule(self);
  19879. }
  19880. template <typename batch_rule_t, batch_rule_t batch_rule>
  19881. ::std::vector<at::Tensor> _foreach_rsqrt_generated_plumbing(at::TensorList self) {
  19882. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19883. auto maybe_layer = maybeCurrentDynamicLayer();
  19884. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19885. int64_t cur_level = maybe_layer->layerId();
  19886. if (!isBatchedAtLevel(self, cur_level)) {
  19887. return at::_ops::_foreach_rsqrt::call(self);
  19888. }
  19889. auto results = batch_rule(self);
  19890. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19891. }
  19892. template <typename batch_rule_t, batch_rule_t batch_rule>
  19893. void _foreach_rsqrt__generated_plumbing(at::TensorList self) {
  19894. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19895. auto maybe_layer = maybeCurrentDynamicLayer();
  19896. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  19897. int64_t cur_level = maybe_layer->layerId();
  19898. if (!isBatchedAtLevel(self, cur_level)) {
  19899. return at::_ops::_foreach_rsqrt_::call(self);
  19900. }
  19901. batch_rule(self);
  19902. }
  19903. template <typename batch_rule_t, batch_rule_t batch_rule>
  19904. ::std::vector<at::Tensor> _foreach_sigmoid_generated_plumbing(at::TensorList self) {
  19905. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19906. auto maybe_layer = maybeCurrentDynamicLayer();
  19907. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19908. int64_t cur_level = maybe_layer->layerId();
  19909. if (!isBatchedAtLevel(self, cur_level)) {
  19910. return at::_ops::_foreach_sigmoid::call(self);
  19911. }
  19912. auto results = batch_rule(self);
  19913. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19914. }
  19915. template <typename batch_rule_t, batch_rule_t batch_rule>
  19916. void _foreach_sigmoid__generated_plumbing(at::TensorList self) {
  19917. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19918. auto maybe_layer = maybeCurrentDynamicLayer();
  19919. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  19920. int64_t cur_level = maybe_layer->layerId();
  19921. if (!isBatchedAtLevel(self, cur_level)) {
  19922. return at::_ops::_foreach_sigmoid_::call(self);
  19923. }
  19924. batch_rule(self);
  19925. }
  19926. template <typename batch_rule_t, batch_rule_t batch_rule>
  19927. ::std::vector<at::Tensor> _foreach_sign_generated_plumbing(at::TensorList self) {
  19928. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19929. auto maybe_layer = maybeCurrentDynamicLayer();
  19930. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19931. int64_t cur_level = maybe_layer->layerId();
  19932. if (!isBatchedAtLevel(self, cur_level)) {
  19933. return at::_ops::_foreach_sign::call(self);
  19934. }
  19935. auto results = batch_rule(self);
  19936. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19937. }
  19938. template <typename batch_rule_t, batch_rule_t batch_rule>
  19939. void _foreach_sign__generated_plumbing(at::TensorList self) {
  19940. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19941. auto maybe_layer = maybeCurrentDynamicLayer();
  19942. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  19943. int64_t cur_level = maybe_layer->layerId();
  19944. if (!isBatchedAtLevel(self, cur_level)) {
  19945. return at::_ops::_foreach_sign_::call(self);
  19946. }
  19947. batch_rule(self);
  19948. }
  19949. template <typename batch_rule_t, batch_rule_t batch_rule>
  19950. ::std::vector<at::Tensor> _foreach_sin_generated_plumbing(at::TensorList self) {
  19951. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19952. auto maybe_layer = maybeCurrentDynamicLayer();
  19953. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19954. int64_t cur_level = maybe_layer->layerId();
  19955. if (!isBatchedAtLevel(self, cur_level)) {
  19956. return at::_ops::_foreach_sin::call(self);
  19957. }
  19958. auto results = batch_rule(self);
  19959. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19960. }
  19961. template <typename batch_rule_t, batch_rule_t batch_rule>
  19962. void _foreach_sin__generated_plumbing(at::TensorList self) {
  19963. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19964. auto maybe_layer = maybeCurrentDynamicLayer();
  19965. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  19966. int64_t cur_level = maybe_layer->layerId();
  19967. if (!isBatchedAtLevel(self, cur_level)) {
  19968. return at::_ops::_foreach_sin_::call(self);
  19969. }
  19970. batch_rule(self);
  19971. }
  19972. template <typename batch_rule_t, batch_rule_t batch_rule>
  19973. ::std::vector<at::Tensor> _foreach_sinh_generated_plumbing(at::TensorList self) {
  19974. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19975. auto maybe_layer = maybeCurrentDynamicLayer();
  19976. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  19977. int64_t cur_level = maybe_layer->layerId();
  19978. if (!isBatchedAtLevel(self, cur_level)) {
  19979. return at::_ops::_foreach_sinh::call(self);
  19980. }
  19981. auto results = batch_rule(self);
  19982. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  19983. }
  19984. template <typename batch_rule_t, batch_rule_t batch_rule>
  19985. void _foreach_sinh__generated_plumbing(at::TensorList self) {
  19986. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19987. auto maybe_layer = maybeCurrentDynamicLayer();
  19988. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  19989. int64_t cur_level = maybe_layer->layerId();
  19990. if (!isBatchedAtLevel(self, cur_level)) {
  19991. return at::_ops::_foreach_sinh_::call(self);
  19992. }
  19993. batch_rule(self);
  19994. }
  19995. template <typename batch_rule_t, batch_rule_t batch_rule>
  19996. ::std::vector<at::Tensor> _foreach_sqrt_generated_plumbing(at::TensorList self) {
  19997. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  19998. auto maybe_layer = maybeCurrentDynamicLayer();
  19999. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20000. int64_t cur_level = maybe_layer->layerId();
  20001. if (!isBatchedAtLevel(self, cur_level)) {
  20002. return at::_ops::_foreach_sqrt::call(self);
  20003. }
  20004. auto results = batch_rule(self);
  20005. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  20006. }
  20007. template <typename batch_rule_t, batch_rule_t batch_rule>
  20008. void _foreach_sqrt__generated_plumbing(at::TensorList self) {
  20009. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20010. auto maybe_layer = maybeCurrentDynamicLayer();
  20011. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  20012. int64_t cur_level = maybe_layer->layerId();
  20013. if (!isBatchedAtLevel(self, cur_level)) {
  20014. return at::_ops::_foreach_sqrt_::call(self);
  20015. }
  20016. batch_rule(self);
  20017. }
  20018. template <typename batch_rule_t, batch_rule_t batch_rule>
  20019. ::std::vector<at::Tensor> _foreach_tan_generated_plumbing(at::TensorList self) {
  20020. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20021. auto maybe_layer = maybeCurrentDynamicLayer();
  20022. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20023. int64_t cur_level = maybe_layer->layerId();
  20024. if (!isBatchedAtLevel(self, cur_level)) {
  20025. return at::_ops::_foreach_tan::call(self);
  20026. }
  20027. auto results = batch_rule(self);
  20028. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  20029. }
  20030. template <typename batch_rule_t, batch_rule_t batch_rule>
  20031. void _foreach_tan__generated_plumbing(at::TensorList self) {
  20032. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20033. auto maybe_layer = maybeCurrentDynamicLayer();
  20034. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  20035. int64_t cur_level = maybe_layer->layerId();
  20036. if (!isBatchedAtLevel(self, cur_level)) {
  20037. return at::_ops::_foreach_tan_::call(self);
  20038. }
  20039. batch_rule(self);
  20040. }
  20041. template <typename batch_rule_t, batch_rule_t batch_rule>
  20042. ::std::vector<at::Tensor> _foreach_tanh_generated_plumbing(at::TensorList self) {
  20043. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20044. auto maybe_layer = maybeCurrentDynamicLayer();
  20045. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20046. int64_t cur_level = maybe_layer->layerId();
  20047. if (!isBatchedAtLevel(self, cur_level)) {
  20048. return at::_ops::_foreach_tanh::call(self);
  20049. }
  20050. auto results = batch_rule(self);
  20051. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  20052. }
  20053. template <typename batch_rule_t, batch_rule_t batch_rule>
  20054. void _foreach_tanh__generated_plumbing(at::TensorList self) {
  20055. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20056. auto maybe_layer = maybeCurrentDynamicLayer();
  20057. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  20058. int64_t cur_level = maybe_layer->layerId();
  20059. if (!isBatchedAtLevel(self, cur_level)) {
  20060. return at::_ops::_foreach_tanh_::call(self);
  20061. }
  20062. batch_rule(self);
  20063. }
  20064. template <typename batch_rule_t, batch_rule_t batch_rule>
  20065. ::std::vector<at::Tensor> _foreach_trunc_generated_plumbing(at::TensorList self) {
  20066. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20067. auto maybe_layer = maybeCurrentDynamicLayer();
  20068. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20069. int64_t cur_level = maybe_layer->layerId();
  20070. if (!isBatchedAtLevel(self, cur_level)) {
  20071. return at::_ops::_foreach_trunc::call(self);
  20072. }
  20073. auto results = batch_rule(self);
  20074. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  20075. }
  20076. template <typename batch_rule_t, batch_rule_t batch_rule>
  20077. void _foreach_trunc__generated_plumbing(at::TensorList self) {
  20078. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20079. auto maybe_layer = maybeCurrentDynamicLayer();
  20080. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  20081. int64_t cur_level = maybe_layer->layerId();
  20082. if (!isBatchedAtLevel(self, cur_level)) {
  20083. return at::_ops::_foreach_trunc_::call(self);
  20084. }
  20085. batch_rule(self);
  20086. }
  20087. template <typename batch_rule_t, batch_rule_t batch_rule>
  20088. void _foreach_zero__generated_plumbing(at::TensorList self) {
  20089. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20090. auto maybe_layer = maybeCurrentDynamicLayer();
  20091. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  20092. int64_t cur_level = maybe_layer->layerId();
  20093. if (!isBatchedAtLevel(self, cur_level)) {
  20094. return at::_ops::_foreach_zero_::call(self);
  20095. }
  20096. batch_rule(self);
  20097. }
  20098. template <typename batch_rule_t, batch_rule_t batch_rule>
  20099. void _foreach_copy__generated_plumbing(at::TensorList self, at::TensorList src, bool non_blocking) {
  20100. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20101. auto maybe_layer = maybeCurrentDynamicLayer();
  20102. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  20103. int64_t cur_level = maybe_layer->layerId();
  20104. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(src, cur_level)) {
  20105. return at::_ops::_foreach_copy_::call(self, src, non_blocking);
  20106. }
  20107. batch_rule(self, src, non_blocking);
  20108. }
  20109. template <typename batch_rule_t, batch_rule_t batch_rule>
  20110. ::std::vector<at::Tensor> _foreach_copy_generated_plumbing(at::TensorList self, at::TensorList src, bool non_blocking) {
  20111. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20112. auto maybe_layer = maybeCurrentDynamicLayer();
  20113. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20114. int64_t cur_level = maybe_layer->layerId();
  20115. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(src, cur_level)) {
  20116. return at::_ops::_foreach_copy::call(self, src, non_blocking);
  20117. }
  20118. auto results = batch_rule(self, src, non_blocking);
  20119. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  20120. }
  20121. template <typename batch_rule_t, batch_rule_t batch_rule>
  20122. at::Tensor bucketize_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & boundaries, bool out_int32, bool right) {
  20123. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20124. auto maybe_layer = maybeCurrentDynamicLayer();
  20125. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20126. int64_t cur_level = maybe_layer->layerId();
  20127. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(boundaries, cur_level)) {
  20128. return at::_ops::bucketize_Tensor::call(self, boundaries, out_int32, right);
  20129. }
  20130. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20131. auto [boundaries_value, boundaries_bdim] = unwrapTensorAtLevel(boundaries, cur_level);
  20132. auto results = batch_rule(self_value, self_bdim, boundaries_value, boundaries_bdim, out_int32, right);
  20133. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20134. }
  20135. template <typename batch_rule_t, batch_rule_t batch_rule>
  20136. at::Tensor bucketize_Scalar_generated_plumbing(const at::Scalar & self, const at::Tensor & boundaries, bool out_int32, bool right) {
  20137. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20138. auto maybe_layer = maybeCurrentDynamicLayer();
  20139. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20140. int64_t cur_level = maybe_layer->layerId();
  20141. if (!isBatchedAtLevel(boundaries, cur_level)) {
  20142. return at::_ops::bucketize_Scalar::call(self, boundaries, out_int32, right);
  20143. }
  20144. auto [boundaries_value, boundaries_bdim] = unwrapTensorAtLevel(boundaries, cur_level);
  20145. auto results = batch_rule(self, boundaries_value, boundaries_bdim, out_int32, right);
  20146. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20147. }
  20148. template <typename batch_rule_t, batch_rule_t batch_rule>
  20149. at::Tensor searchsorted_Tensor_generated_plumbing(const at::Tensor & sorted_sequence, const at::Tensor & self, bool out_int32, bool right, ::std::optional<c10::string_view> side, const ::std::optional<at::Tensor> & sorter) {
  20150. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20151. auto maybe_layer = maybeCurrentDynamicLayer();
  20152. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20153. int64_t cur_level = maybe_layer->layerId();
  20154. if (!isBatchedAtLevel(sorted_sequence, cur_level) && !isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(sorter, cur_level)) {
  20155. return at::_ops::searchsorted_Tensor::call(sorted_sequence, self, out_int32, right, side, sorter);
  20156. }
  20157. auto [sorted_sequence_value, sorted_sequence_bdim] = unwrapTensorAtLevel(sorted_sequence, cur_level);
  20158. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20159. std::optional<Tensor> sorter_value;
  20160. std::optional<int64_t> sorter_bdim;
  20161. if (sorter) {
  20162. std::tie(sorter_value, sorter_bdim) = unwrapTensorAtLevel(sorter.value(), cur_level);
  20163. }
  20164. auto results = batch_rule(sorted_sequence_value, sorted_sequence_bdim, self_value, self_bdim, out_int32, right, side, sorter_value, sorter_bdim);
  20165. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20166. }
  20167. template <typename batch_rule_t, batch_rule_t batch_rule>
  20168. at::Tensor searchsorted_Scalar_generated_plumbing(const at::Tensor & sorted_sequence, const at::Scalar & self, bool out_int32, bool right, ::std::optional<c10::string_view> side, const ::std::optional<at::Tensor> & sorter) {
  20169. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20170. auto maybe_layer = maybeCurrentDynamicLayer();
  20171. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20172. int64_t cur_level = maybe_layer->layerId();
  20173. if (!isBatchedAtLevel(sorted_sequence, cur_level) && !isBatchedAtLevel(sorter, cur_level)) {
  20174. return at::_ops::searchsorted_Scalar::call(sorted_sequence, self, out_int32, right, side, sorter);
  20175. }
  20176. auto [sorted_sequence_value, sorted_sequence_bdim] = unwrapTensorAtLevel(sorted_sequence, cur_level);
  20177. std::optional<Tensor> sorter_value;
  20178. std::optional<int64_t> sorter_bdim;
  20179. if (sorter) {
  20180. std::tie(sorter_value, sorter_bdim) = unwrapTensorAtLevel(sorter.value(), cur_level);
  20181. }
  20182. auto results = batch_rule(sorted_sequence_value, sorted_sequence_bdim, self, out_int32, right, side, sorter_value, sorter_bdim);
  20183. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20184. }
  20185. template <typename batch_rule_t, batch_rule_t batch_rule>
  20186. at::Tensor _convert_indices_from_coo_to_csr_generated_plumbing(const at::Tensor & self, int64_t size, bool out_int32) {
  20187. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20188. auto maybe_layer = maybeCurrentDynamicLayer();
  20189. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20190. int64_t cur_level = maybe_layer->layerId();
  20191. if (!isBatchedAtLevel(self, cur_level)) {
  20192. return at::_ops::_convert_indices_from_coo_to_csr::call(self, size, out_int32);
  20193. }
  20194. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20195. auto results = batch_rule(self_value, self_bdim, size, out_int32);
  20196. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20197. }
  20198. template <typename batch_rule_t, batch_rule_t batch_rule>
  20199. at::Tensor _convert_indices_from_csr_to_coo_generated_plumbing(const at::Tensor & crow_indices, const at::Tensor & col_indices, bool out_int32, bool transpose) {
  20200. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20201. auto maybe_layer = maybeCurrentDynamicLayer();
  20202. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20203. int64_t cur_level = maybe_layer->layerId();
  20204. if (!isBatchedAtLevel(crow_indices, cur_level) && !isBatchedAtLevel(col_indices, cur_level)) {
  20205. return at::_ops::_convert_indices_from_csr_to_coo::call(crow_indices, col_indices, out_int32, transpose);
  20206. }
  20207. auto [crow_indices_value, crow_indices_bdim] = unwrapTensorAtLevel(crow_indices, cur_level);
  20208. auto [col_indices_value, col_indices_bdim] = unwrapTensorAtLevel(col_indices, cur_level);
  20209. auto results = batch_rule(crow_indices_value, crow_indices_bdim, col_indices_value, col_indices_bdim, out_int32, transpose);
  20210. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20211. }
  20212. template <typename batch_rule_t, batch_rule_t batch_rule>
  20213. at::Tensor mse_loss_generated_plumbing(const at::Tensor & self, const at::Tensor & target, int64_t reduction) {
  20214. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20215. auto maybe_layer = maybeCurrentDynamicLayer();
  20216. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20217. int64_t cur_level = maybe_layer->layerId();
  20218. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(target, cur_level)) {
  20219. return at::_ops::mse_loss::call(self, target, reduction);
  20220. }
  20221. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20222. auto [target_value, target_bdim] = unwrapTensorAtLevel(target, cur_level);
  20223. auto results = batch_rule(self_value, self_bdim, target_value, target_bdim, reduction);
  20224. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20225. }
  20226. template <typename batch_rule_t, batch_rule_t batch_rule>
  20227. at::Tensor mse_loss_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self, const at::Tensor & target, int64_t reduction) {
  20228. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20229. auto maybe_layer = maybeCurrentDynamicLayer();
  20230. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20231. int64_t cur_level = maybe_layer->layerId();
  20232. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(target, cur_level)) {
  20233. return at::_ops::mse_loss_backward::call(grad_output, self, target, reduction);
  20234. }
  20235. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  20236. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20237. auto [target_value, target_bdim] = unwrapTensorAtLevel(target, cur_level);
  20238. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim, target_value, target_bdim, reduction);
  20239. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20240. }
  20241. template <typename batch_rule_t, batch_rule_t batch_rule>
  20242. at::Tensor l1_loss_generated_plumbing(const at::Tensor & self, const at::Tensor & target, int64_t reduction) {
  20243. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20244. auto maybe_layer = maybeCurrentDynamicLayer();
  20245. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20246. int64_t cur_level = maybe_layer->layerId();
  20247. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(target, cur_level)) {
  20248. return at::_ops::l1_loss::call(self, target, reduction);
  20249. }
  20250. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20251. auto [target_value, target_bdim] = unwrapTensorAtLevel(target, cur_level);
  20252. auto results = batch_rule(self_value, self_bdim, target_value, target_bdim, reduction);
  20253. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20254. }
  20255. template <typename batch_rule_t, batch_rule_t batch_rule>
  20256. at::Tensor multi_margin_loss_generated_plumbing(const at::Tensor & self, const at::Tensor & target, const at::Scalar & p, const at::Scalar & margin, const ::std::optional<at::Tensor> & weight, int64_t reduction) {
  20257. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20258. auto maybe_layer = maybeCurrentDynamicLayer();
  20259. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20260. int64_t cur_level = maybe_layer->layerId();
  20261. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(target, cur_level) && !isBatchedAtLevel(weight, cur_level)) {
  20262. return at::_ops::multi_margin_loss::call(self, target, p, margin, weight, reduction);
  20263. }
  20264. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20265. auto [target_value, target_bdim] = unwrapTensorAtLevel(target, cur_level);
  20266. std::optional<Tensor> weight_value;
  20267. std::optional<int64_t> weight_bdim;
  20268. if (weight) {
  20269. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  20270. }
  20271. auto results = batch_rule(self_value, self_bdim, target_value, target_bdim, p, margin, weight_value, weight_bdim, reduction);
  20272. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20273. }
  20274. template <typename batch_rule_t, batch_rule_t batch_rule>
  20275. at::Tensor multi_margin_loss_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self, const at::Tensor & target, const at::Scalar & p, const at::Scalar & margin, const ::std::optional<at::Tensor> & weight, int64_t reduction) {
  20276. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20277. auto maybe_layer = maybeCurrentDynamicLayer();
  20278. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20279. int64_t cur_level = maybe_layer->layerId();
  20280. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(target, cur_level) && !isBatchedAtLevel(weight, cur_level)) {
  20281. return at::_ops::multi_margin_loss_backward::call(grad_output, self, target, p, margin, weight, reduction);
  20282. }
  20283. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  20284. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20285. auto [target_value, target_bdim] = unwrapTensorAtLevel(target, cur_level);
  20286. std::optional<Tensor> weight_value;
  20287. std::optional<int64_t> weight_bdim;
  20288. if (weight) {
  20289. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  20290. }
  20291. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim, target_value, target_bdim, p, margin, weight_value, weight_bdim, reduction);
  20292. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20293. }
  20294. template <typename batch_rule_t, batch_rule_t batch_rule>
  20295. at::Tensor multilabel_margin_loss_generated_plumbing(const at::Tensor & self, const at::Tensor & target, int64_t reduction) {
  20296. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20297. auto maybe_layer = maybeCurrentDynamicLayer();
  20298. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20299. int64_t cur_level = maybe_layer->layerId();
  20300. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(target, cur_level)) {
  20301. return at::_ops::multilabel_margin_loss::call(self, target, reduction);
  20302. }
  20303. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20304. auto [target_value, target_bdim] = unwrapTensorAtLevel(target, cur_level);
  20305. auto results = batch_rule(self_value, self_bdim, target_value, target_bdim, reduction);
  20306. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20307. }
  20308. template <typename batch_rule_t, batch_rule_t batch_rule>
  20309. ::std::tuple<at::Tensor,at::Tensor> multilabel_margin_loss_forward_generated_plumbing(const at::Tensor & self, const at::Tensor & target, int64_t reduction) {
  20310. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20311. auto maybe_layer = maybeCurrentDynamicLayer();
  20312. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20313. int64_t cur_level = maybe_layer->layerId();
  20314. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(target, cur_level)) {
  20315. return at::_ops::multilabel_margin_loss_forward::call(self, target, reduction);
  20316. }
  20317. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20318. auto [target_value, target_bdim] = unwrapTensorAtLevel(target, cur_level);
  20319. auto results = batch_rule(self_value, self_bdim, target_value, target_bdim, reduction);
  20320. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  20321. }
  20322. template <typename batch_rule_t, batch_rule_t batch_rule>
  20323. at::Tensor multilabel_margin_loss_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self, const at::Tensor & target, int64_t reduction, const at::Tensor & is_target) {
  20324. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20325. auto maybe_layer = maybeCurrentDynamicLayer();
  20326. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20327. int64_t cur_level = maybe_layer->layerId();
  20328. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(target, cur_level) && !isBatchedAtLevel(is_target, cur_level)) {
  20329. return at::_ops::multilabel_margin_loss_backward::call(grad_output, self, target, reduction, is_target);
  20330. }
  20331. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  20332. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20333. auto [target_value, target_bdim] = unwrapTensorAtLevel(target, cur_level);
  20334. auto [is_target_value, is_target_bdim] = unwrapTensorAtLevel(is_target, cur_level);
  20335. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim, target_value, target_bdim, reduction, is_target_value, is_target_bdim);
  20336. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20337. }
  20338. template <typename batch_rule_t, batch_rule_t batch_rule>
  20339. at::Tensor nll_loss_nd_generated_plumbing(const at::Tensor & self, const at::Tensor & target, const ::std::optional<at::Tensor> & weight, int64_t reduction, c10::SymInt ignore_index) {
  20340. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20341. auto maybe_layer = maybeCurrentDynamicLayer();
  20342. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20343. int64_t cur_level = maybe_layer->layerId();
  20344. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(target, cur_level) && !isBatchedAtLevel(weight, cur_level)) {
  20345. return at::_ops::nll_loss_nd::call(self, target, weight, reduction, ignore_index);
  20346. }
  20347. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20348. auto [target_value, target_bdim] = unwrapTensorAtLevel(target, cur_level);
  20349. std::optional<Tensor> weight_value;
  20350. std::optional<int64_t> weight_bdim;
  20351. if (weight) {
  20352. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  20353. }
  20354. auto results = batch_rule(self_value, self_bdim, target_value, target_bdim, weight_value, weight_bdim, reduction, ignore_index);
  20355. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20356. }
  20357. template <typename batch_rule_t, batch_rule_t batch_rule>
  20358. at::Tensor nll_loss_generated_plumbing(const at::Tensor & self, const at::Tensor & target, const ::std::optional<at::Tensor> & weight, int64_t reduction, c10::SymInt ignore_index) {
  20359. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20360. auto maybe_layer = maybeCurrentDynamicLayer();
  20361. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20362. int64_t cur_level = maybe_layer->layerId();
  20363. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(target, cur_level) && !isBatchedAtLevel(weight, cur_level)) {
  20364. return at::_ops::nll_loss::call(self, target, weight, reduction, ignore_index);
  20365. }
  20366. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20367. auto [target_value, target_bdim] = unwrapTensorAtLevel(target, cur_level);
  20368. std::optional<Tensor> weight_value;
  20369. std::optional<int64_t> weight_bdim;
  20370. if (weight) {
  20371. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  20372. }
  20373. auto results = batch_rule(self_value, self_bdim, target_value, target_bdim, weight_value, weight_bdim, reduction, ignore_index);
  20374. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20375. }
  20376. template <typename batch_rule_t, batch_rule_t batch_rule>
  20377. ::std::tuple<at::Tensor,at::Tensor> nll_loss_forward_generated_plumbing(const at::Tensor & self, const at::Tensor & target, const ::std::optional<at::Tensor> & weight, int64_t reduction, c10::SymInt ignore_index) {
  20378. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20379. auto maybe_layer = maybeCurrentDynamicLayer();
  20380. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20381. int64_t cur_level = maybe_layer->layerId();
  20382. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(target, cur_level) && !isBatchedAtLevel(weight, cur_level)) {
  20383. return at::_ops::nll_loss_forward::call(self, target, weight, reduction, ignore_index);
  20384. }
  20385. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20386. auto [target_value, target_bdim] = unwrapTensorAtLevel(target, cur_level);
  20387. std::optional<Tensor> weight_value;
  20388. std::optional<int64_t> weight_bdim;
  20389. if (weight) {
  20390. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  20391. }
  20392. auto results = batch_rule(self_value, self_bdim, target_value, target_bdim, weight_value, weight_bdim, reduction, ignore_index);
  20393. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  20394. }
  20395. template <typename batch_rule_t, batch_rule_t batch_rule>
  20396. at::Tensor nll_loss_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self, const at::Tensor & target, const ::std::optional<at::Tensor> & weight, int64_t reduction, c10::SymInt ignore_index, const at::Tensor & total_weight) {
  20397. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20398. auto maybe_layer = maybeCurrentDynamicLayer();
  20399. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20400. int64_t cur_level = maybe_layer->layerId();
  20401. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(target, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(total_weight, cur_level)) {
  20402. return at::_ops::nll_loss_backward::call(grad_output, self, target, weight, reduction, ignore_index, total_weight);
  20403. }
  20404. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  20405. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20406. auto [target_value, target_bdim] = unwrapTensorAtLevel(target, cur_level);
  20407. auto [total_weight_value, total_weight_bdim] = unwrapTensorAtLevel(total_weight, cur_level);
  20408. std::optional<Tensor> weight_value;
  20409. std::optional<int64_t> weight_bdim;
  20410. if (weight) {
  20411. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  20412. }
  20413. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim, target_value, target_bdim, weight_value, weight_bdim, reduction, ignore_index, total_weight_value, total_weight_bdim);
  20414. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20415. }
  20416. template <typename batch_rule_t, batch_rule_t batch_rule>
  20417. at::Tensor nll_loss2d_generated_plumbing(const at::Tensor & self, const at::Tensor & target, const ::std::optional<at::Tensor> & weight, int64_t reduction, c10::SymInt ignore_index) {
  20418. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20419. auto maybe_layer = maybeCurrentDynamicLayer();
  20420. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20421. int64_t cur_level = maybe_layer->layerId();
  20422. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(target, cur_level) && !isBatchedAtLevel(weight, cur_level)) {
  20423. return at::_ops::nll_loss2d::call(self, target, weight, reduction, ignore_index);
  20424. }
  20425. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20426. auto [target_value, target_bdim] = unwrapTensorAtLevel(target, cur_level);
  20427. std::optional<Tensor> weight_value;
  20428. std::optional<int64_t> weight_bdim;
  20429. if (weight) {
  20430. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  20431. }
  20432. auto results = batch_rule(self_value, self_bdim, target_value, target_bdim, weight_value, weight_bdim, reduction, ignore_index);
  20433. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20434. }
  20435. template <typename batch_rule_t, batch_rule_t batch_rule>
  20436. ::std::tuple<at::Tensor,at::Tensor> nll_loss2d_forward_generated_plumbing(const at::Tensor & self, const at::Tensor & target, const ::std::optional<at::Tensor> & weight, int64_t reduction, c10::SymInt ignore_index) {
  20437. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20438. auto maybe_layer = maybeCurrentDynamicLayer();
  20439. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20440. int64_t cur_level = maybe_layer->layerId();
  20441. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(target, cur_level) && !isBatchedAtLevel(weight, cur_level)) {
  20442. return at::_ops::nll_loss2d_forward::call(self, target, weight, reduction, ignore_index);
  20443. }
  20444. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20445. auto [target_value, target_bdim] = unwrapTensorAtLevel(target, cur_level);
  20446. std::optional<Tensor> weight_value;
  20447. std::optional<int64_t> weight_bdim;
  20448. if (weight) {
  20449. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  20450. }
  20451. auto results = batch_rule(self_value, self_bdim, target_value, target_bdim, weight_value, weight_bdim, reduction, ignore_index);
  20452. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  20453. }
  20454. template <typename batch_rule_t, batch_rule_t batch_rule>
  20455. at::Tensor nll_loss2d_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self, const at::Tensor & target, const ::std::optional<at::Tensor> & weight, int64_t reduction, c10::SymInt ignore_index, const at::Tensor & total_weight) {
  20456. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20457. auto maybe_layer = maybeCurrentDynamicLayer();
  20458. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20459. int64_t cur_level = maybe_layer->layerId();
  20460. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(target, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(total_weight, cur_level)) {
  20461. return at::_ops::nll_loss2d_backward::call(grad_output, self, target, weight, reduction, ignore_index, total_weight);
  20462. }
  20463. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  20464. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20465. auto [target_value, target_bdim] = unwrapTensorAtLevel(target, cur_level);
  20466. auto [total_weight_value, total_weight_bdim] = unwrapTensorAtLevel(total_weight, cur_level);
  20467. std::optional<Tensor> weight_value;
  20468. std::optional<int64_t> weight_bdim;
  20469. if (weight) {
  20470. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  20471. }
  20472. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim, target_value, target_bdim, weight_value, weight_bdim, reduction, ignore_index, total_weight_value, total_weight_bdim);
  20473. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20474. }
  20475. template <typename batch_rule_t, batch_rule_t batch_rule>
  20476. at::Tensor smooth_l1_loss_generated_plumbing(const at::Tensor & self, const at::Tensor & target, int64_t reduction, double beta) {
  20477. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20478. auto maybe_layer = maybeCurrentDynamicLayer();
  20479. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20480. int64_t cur_level = maybe_layer->layerId();
  20481. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(target, cur_level)) {
  20482. return at::_ops::smooth_l1_loss::call(self, target, reduction, beta);
  20483. }
  20484. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20485. auto [target_value, target_bdim] = unwrapTensorAtLevel(target, cur_level);
  20486. auto results = batch_rule(self_value, self_bdim, target_value, target_bdim, reduction, beta);
  20487. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20488. }
  20489. template <typename batch_rule_t, batch_rule_t batch_rule>
  20490. at::Tensor smooth_l1_loss_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self, const at::Tensor & target, int64_t reduction, double beta) {
  20491. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20492. auto maybe_layer = maybeCurrentDynamicLayer();
  20493. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20494. int64_t cur_level = maybe_layer->layerId();
  20495. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(target, cur_level)) {
  20496. return at::_ops::smooth_l1_loss_backward::call(grad_output, self, target, reduction, beta);
  20497. }
  20498. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  20499. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20500. auto [target_value, target_bdim] = unwrapTensorAtLevel(target, cur_level);
  20501. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim, target_value, target_bdim, reduction, beta);
  20502. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20503. }
  20504. template <typename batch_rule_t, batch_rule_t batch_rule>
  20505. at::Tensor huber_loss_generated_plumbing(const at::Tensor & self, const at::Tensor & target, int64_t reduction, double delta) {
  20506. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20507. auto maybe_layer = maybeCurrentDynamicLayer();
  20508. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20509. int64_t cur_level = maybe_layer->layerId();
  20510. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(target, cur_level)) {
  20511. return at::_ops::huber_loss::call(self, target, reduction, delta);
  20512. }
  20513. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20514. auto [target_value, target_bdim] = unwrapTensorAtLevel(target, cur_level);
  20515. auto results = batch_rule(self_value, self_bdim, target_value, target_bdim, reduction, delta);
  20516. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20517. }
  20518. template <typename batch_rule_t, batch_rule_t batch_rule>
  20519. at::Tensor huber_loss_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self, const at::Tensor & target, int64_t reduction, double delta) {
  20520. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20521. auto maybe_layer = maybeCurrentDynamicLayer();
  20522. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20523. int64_t cur_level = maybe_layer->layerId();
  20524. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(target, cur_level)) {
  20525. return at::_ops::huber_loss_backward::call(grad_output, self, target, reduction, delta);
  20526. }
  20527. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  20528. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20529. auto [target_value, target_bdim] = unwrapTensorAtLevel(target, cur_level);
  20530. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim, target_value, target_bdim, reduction, delta);
  20531. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20532. }
  20533. template <typename batch_rule_t, batch_rule_t batch_rule>
  20534. at::Tensor soft_margin_loss_generated_plumbing(const at::Tensor & self, const at::Tensor & target, int64_t reduction) {
  20535. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20536. auto maybe_layer = maybeCurrentDynamicLayer();
  20537. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20538. int64_t cur_level = maybe_layer->layerId();
  20539. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(target, cur_level)) {
  20540. return at::_ops::soft_margin_loss::call(self, target, reduction);
  20541. }
  20542. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20543. auto [target_value, target_bdim] = unwrapTensorAtLevel(target, cur_level);
  20544. auto results = batch_rule(self_value, self_bdim, target_value, target_bdim, reduction);
  20545. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20546. }
  20547. template <typename batch_rule_t, batch_rule_t batch_rule>
  20548. at::Tensor soft_margin_loss_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self, const at::Tensor & target, int64_t reduction) {
  20549. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20550. auto maybe_layer = maybeCurrentDynamicLayer();
  20551. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20552. int64_t cur_level = maybe_layer->layerId();
  20553. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(target, cur_level)) {
  20554. return at::_ops::soft_margin_loss_backward::call(grad_output, self, target, reduction);
  20555. }
  20556. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  20557. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20558. auto [target_value, target_bdim] = unwrapTensorAtLevel(target, cur_level);
  20559. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim, target_value, target_bdim, reduction);
  20560. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20561. }
  20562. template <typename batch_rule_t, batch_rule_t batch_rule>
  20563. at::Tensor elu_generated_plumbing(const at::Tensor & self, const at::Scalar & alpha, const at::Scalar & scale, const at::Scalar & input_scale) {
  20564. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20565. auto maybe_layer = maybeCurrentDynamicLayer();
  20566. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20567. int64_t cur_level = maybe_layer->layerId();
  20568. if (!isBatchedAtLevel(self, cur_level)) {
  20569. return at::_ops::elu::call(self, alpha, scale, input_scale);
  20570. }
  20571. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20572. auto results = batch_rule(self_value, self_bdim, alpha, scale, input_scale);
  20573. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20574. }
  20575. template <typename batch_rule_t, batch_rule_t batch_rule>
  20576. at::Tensor elu_backward_generated_plumbing(const at::Tensor & grad_output, const at::Scalar & alpha, const at::Scalar & scale, const at::Scalar & input_scale, bool is_result, const at::Tensor & self_or_result) {
  20577. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20578. auto maybe_layer = maybeCurrentDynamicLayer();
  20579. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20580. int64_t cur_level = maybe_layer->layerId();
  20581. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self_or_result, cur_level)) {
  20582. return at::_ops::elu_backward::call(grad_output, alpha, scale, input_scale, is_result, self_or_result);
  20583. }
  20584. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  20585. auto [self_or_result_value, self_or_result_bdim] = unwrapTensorAtLevel(self_or_result, cur_level);
  20586. auto results = batch_rule(grad_output_value, grad_output_bdim, alpha, scale, input_scale, is_result, self_or_result_value, self_or_result_bdim);
  20587. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20588. }
  20589. template <typename batch_rule_t, batch_rule_t batch_rule>
  20590. at::Tensor & elu__generated_plumbing(at::Tensor & self, const at::Scalar & alpha, const at::Scalar & scale, const at::Scalar & input_scale) {
  20591. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20592. auto maybe_layer = maybeCurrentDynamicLayer();
  20593. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  20594. int64_t cur_level = maybe_layer->layerId();
  20595. if (!isBatchedAtLevel(self, cur_level)) {
  20596. return at::_ops::elu_::call(self, alpha, scale, input_scale);
  20597. }
  20598. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20599. batch_rule(self_value, self_bdim, alpha, scale, input_scale);
  20600. return self;
  20601. }
  20602. template <typename batch_rule_t, batch_rule_t batch_rule>
  20603. at::Tensor glu_generated_plumbing(const at::Tensor & self, int64_t dim) {
  20604. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20605. auto maybe_layer = maybeCurrentDynamicLayer();
  20606. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20607. int64_t cur_level = maybe_layer->layerId();
  20608. if (!isBatchedAtLevel(self, cur_level)) {
  20609. return at::_ops::glu::call(self, dim);
  20610. }
  20611. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20612. auto results = batch_rule(self_value, self_bdim, dim);
  20613. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20614. }
  20615. template <typename batch_rule_t, batch_rule_t batch_rule>
  20616. at::Tensor glu_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self, int64_t dim) {
  20617. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20618. auto maybe_layer = maybeCurrentDynamicLayer();
  20619. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20620. int64_t cur_level = maybe_layer->layerId();
  20621. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level)) {
  20622. return at::_ops::glu_backward::call(grad_output, self, dim);
  20623. }
  20624. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  20625. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20626. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim, dim);
  20627. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20628. }
  20629. template <typename batch_rule_t, batch_rule_t batch_rule>
  20630. at::Tensor glu_jvp_generated_plumbing(const at::Tensor & glu, const at::Tensor & x, const at::Tensor & dx, int64_t dim) {
  20631. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20632. auto maybe_layer = maybeCurrentDynamicLayer();
  20633. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20634. int64_t cur_level = maybe_layer->layerId();
  20635. if (!isBatchedAtLevel(glu, cur_level) && !isBatchedAtLevel(x, cur_level) && !isBatchedAtLevel(dx, cur_level)) {
  20636. return at::_ops::glu_jvp::call(glu, x, dx, dim);
  20637. }
  20638. auto [glu_value, glu_bdim] = unwrapTensorAtLevel(glu, cur_level);
  20639. auto [x_value, x_bdim] = unwrapTensorAtLevel(x, cur_level);
  20640. auto [dx_value, dx_bdim] = unwrapTensorAtLevel(dx, cur_level);
  20641. auto results = batch_rule(glu_value, glu_bdim, x_value, x_bdim, dx_value, dx_bdim, dim);
  20642. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20643. }
  20644. template <typename batch_rule_t, batch_rule_t batch_rule>
  20645. at::Tensor glu_backward_jvp_generated_plumbing(const at::Tensor & grad_x, const at::Tensor & grad_glu, const at::Tensor & x, const at::Tensor & dgrad_glu, const at::Tensor & dx, int64_t dim) {
  20646. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20647. auto maybe_layer = maybeCurrentDynamicLayer();
  20648. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20649. int64_t cur_level = maybe_layer->layerId();
  20650. if (!isBatchedAtLevel(grad_x, cur_level) && !isBatchedAtLevel(grad_glu, cur_level) && !isBatchedAtLevel(x, cur_level) && !isBatchedAtLevel(dgrad_glu, cur_level) && !isBatchedAtLevel(dx, cur_level)) {
  20651. return at::_ops::glu_backward_jvp::call(grad_x, grad_glu, x, dgrad_glu, dx, dim);
  20652. }
  20653. auto [grad_x_value, grad_x_bdim] = unwrapTensorAtLevel(grad_x, cur_level);
  20654. auto [grad_glu_value, grad_glu_bdim] = unwrapTensorAtLevel(grad_glu, cur_level);
  20655. auto [x_value, x_bdim] = unwrapTensorAtLevel(x, cur_level);
  20656. auto [dgrad_glu_value, dgrad_glu_bdim] = unwrapTensorAtLevel(dgrad_glu, cur_level);
  20657. auto [dx_value, dx_bdim] = unwrapTensorAtLevel(dx, cur_level);
  20658. auto results = batch_rule(grad_x_value, grad_x_bdim, grad_glu_value, grad_glu_bdim, x_value, x_bdim, dgrad_glu_value, dgrad_glu_bdim, dx_value, dx_bdim, dim);
  20659. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20660. }
  20661. template <typename batch_rule_t, batch_rule_t batch_rule>
  20662. at::Tensor hardsigmoid_generated_plumbing(const at::Tensor & self) {
  20663. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20664. auto maybe_layer = maybeCurrentDynamicLayer();
  20665. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20666. int64_t cur_level = maybe_layer->layerId();
  20667. if (!isBatchedAtLevel(self, cur_level)) {
  20668. return at::_ops::hardsigmoid::call(self);
  20669. }
  20670. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20671. auto results = batch_rule(self_value, self_bdim);
  20672. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20673. }
  20674. template <typename batch_rule_t, batch_rule_t batch_rule>
  20675. at::Tensor & hardsigmoid__generated_plumbing(at::Tensor & self) {
  20676. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20677. auto maybe_layer = maybeCurrentDynamicLayer();
  20678. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  20679. int64_t cur_level = maybe_layer->layerId();
  20680. if (!isBatchedAtLevel(self, cur_level)) {
  20681. return at::_ops::hardsigmoid_::call(self);
  20682. }
  20683. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20684. batch_rule(self_value, self_bdim);
  20685. return self;
  20686. }
  20687. template <typename batch_rule_t, batch_rule_t batch_rule>
  20688. at::Tensor hardsigmoid_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self) {
  20689. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20690. auto maybe_layer = maybeCurrentDynamicLayer();
  20691. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20692. int64_t cur_level = maybe_layer->layerId();
  20693. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level)) {
  20694. return at::_ops::hardsigmoid_backward::call(grad_output, self);
  20695. }
  20696. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  20697. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20698. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim);
  20699. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20700. }
  20701. template <typename batch_rule_t, batch_rule_t batch_rule>
  20702. at::Tensor hardtanh_generated_plumbing(const at::Tensor & self, const at::Scalar & min_val, const at::Scalar & max_val) {
  20703. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20704. auto maybe_layer = maybeCurrentDynamicLayer();
  20705. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20706. int64_t cur_level = maybe_layer->layerId();
  20707. if (!isBatchedAtLevel(self, cur_level)) {
  20708. return at::_ops::hardtanh::call(self, min_val, max_val);
  20709. }
  20710. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20711. auto results = batch_rule(self_value, self_bdim, min_val, max_val);
  20712. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20713. }
  20714. template <typename batch_rule_t, batch_rule_t batch_rule>
  20715. at::Tensor hardtanh_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self, const at::Scalar & min_val, const at::Scalar & max_val) {
  20716. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20717. auto maybe_layer = maybeCurrentDynamicLayer();
  20718. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20719. int64_t cur_level = maybe_layer->layerId();
  20720. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level)) {
  20721. return at::_ops::hardtanh_backward::call(grad_output, self, min_val, max_val);
  20722. }
  20723. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  20724. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20725. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim, min_val, max_val);
  20726. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20727. }
  20728. template <typename batch_rule_t, batch_rule_t batch_rule>
  20729. at::Tensor & hardtanh__generated_plumbing(at::Tensor & self, const at::Scalar & min_val, const at::Scalar & max_val) {
  20730. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20731. auto maybe_layer = maybeCurrentDynamicLayer();
  20732. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  20733. int64_t cur_level = maybe_layer->layerId();
  20734. if (!isBatchedAtLevel(self, cur_level)) {
  20735. return at::_ops::hardtanh_::call(self, min_val, max_val);
  20736. }
  20737. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20738. batch_rule(self_value, self_bdim, min_val, max_val);
  20739. return self;
  20740. }
  20741. template <typename batch_rule_t, batch_rule_t batch_rule>
  20742. at::Tensor hardswish_generated_plumbing(const at::Tensor & self) {
  20743. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20744. auto maybe_layer = maybeCurrentDynamicLayer();
  20745. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20746. int64_t cur_level = maybe_layer->layerId();
  20747. if (!isBatchedAtLevel(self, cur_level)) {
  20748. return at::_ops::hardswish::call(self);
  20749. }
  20750. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20751. auto results = batch_rule(self_value, self_bdim);
  20752. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20753. }
  20754. template <typename batch_rule_t, batch_rule_t batch_rule>
  20755. at::Tensor & hardswish__generated_plumbing(at::Tensor & self) {
  20756. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20757. auto maybe_layer = maybeCurrentDynamicLayer();
  20758. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  20759. int64_t cur_level = maybe_layer->layerId();
  20760. if (!isBatchedAtLevel(self, cur_level)) {
  20761. return at::_ops::hardswish_::call(self);
  20762. }
  20763. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20764. batch_rule(self_value, self_bdim);
  20765. return self;
  20766. }
  20767. template <typename batch_rule_t, batch_rule_t batch_rule>
  20768. at::Tensor hardswish_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self) {
  20769. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20770. auto maybe_layer = maybeCurrentDynamicLayer();
  20771. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20772. int64_t cur_level = maybe_layer->layerId();
  20773. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level)) {
  20774. return at::_ops::hardswish_backward::call(grad_output, self);
  20775. }
  20776. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  20777. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20778. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim);
  20779. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20780. }
  20781. template <typename batch_rule_t, batch_rule_t batch_rule>
  20782. at::Tensor leaky_relu_generated_plumbing(const at::Tensor & self, const at::Scalar & negative_slope) {
  20783. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20784. auto maybe_layer = maybeCurrentDynamicLayer();
  20785. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20786. int64_t cur_level = maybe_layer->layerId();
  20787. if (!isBatchedAtLevel(self, cur_level)) {
  20788. return at::_ops::leaky_relu::call(self, negative_slope);
  20789. }
  20790. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20791. auto results = batch_rule(self_value, self_bdim, negative_slope);
  20792. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20793. }
  20794. template <typename batch_rule_t, batch_rule_t batch_rule>
  20795. at::Tensor leaky_relu_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self, const at::Scalar & negative_slope, bool self_is_result) {
  20796. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20797. auto maybe_layer = maybeCurrentDynamicLayer();
  20798. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20799. int64_t cur_level = maybe_layer->layerId();
  20800. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level)) {
  20801. return at::_ops::leaky_relu_backward::call(grad_output, self, negative_slope, self_is_result);
  20802. }
  20803. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  20804. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20805. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim, negative_slope, self_is_result);
  20806. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20807. }
  20808. template <typename batch_rule_t, batch_rule_t batch_rule>
  20809. at::Tensor & leaky_relu__generated_plumbing(at::Tensor & self, const at::Scalar & negative_slope) {
  20810. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20811. auto maybe_layer = maybeCurrentDynamicLayer();
  20812. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  20813. int64_t cur_level = maybe_layer->layerId();
  20814. if (!isBatchedAtLevel(self, cur_level)) {
  20815. return at::_ops::leaky_relu_::call(self, negative_slope);
  20816. }
  20817. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20818. batch_rule(self_value, self_bdim, negative_slope);
  20819. return self;
  20820. }
  20821. template <typename batch_rule_t, batch_rule_t batch_rule>
  20822. at::Tensor log_sigmoid_generated_plumbing(const at::Tensor & self) {
  20823. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20824. auto maybe_layer = maybeCurrentDynamicLayer();
  20825. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20826. int64_t cur_level = maybe_layer->layerId();
  20827. if (!isBatchedAtLevel(self, cur_level)) {
  20828. return at::_ops::log_sigmoid::call(self);
  20829. }
  20830. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20831. auto results = batch_rule(self_value, self_bdim);
  20832. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20833. }
  20834. template <typename batch_rule_t, batch_rule_t batch_rule>
  20835. ::std::tuple<at::Tensor,at::Tensor> log_sigmoid_forward_generated_plumbing(const at::Tensor & self) {
  20836. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20837. auto maybe_layer = maybeCurrentDynamicLayer();
  20838. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20839. int64_t cur_level = maybe_layer->layerId();
  20840. if (!isBatchedAtLevel(self, cur_level)) {
  20841. return at::_ops::log_sigmoid_forward::call(self);
  20842. }
  20843. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20844. auto results = batch_rule(self_value, self_bdim);
  20845. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  20846. }
  20847. template <typename batch_rule_t, batch_rule_t batch_rule>
  20848. at::Tensor log_sigmoid_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self, const at::Tensor & buffer) {
  20849. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20850. auto maybe_layer = maybeCurrentDynamicLayer();
  20851. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20852. int64_t cur_level = maybe_layer->layerId();
  20853. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(buffer, cur_level)) {
  20854. return at::_ops::log_sigmoid_backward::call(grad_output, self, buffer);
  20855. }
  20856. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  20857. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20858. auto [buffer_value, buffer_bdim] = unwrapTensorAtLevel(buffer, cur_level);
  20859. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim, buffer_value, buffer_bdim);
  20860. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20861. }
  20862. template <typename batch_rule_t, batch_rule_t batch_rule>
  20863. at::Tensor rrelu_with_noise_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self, const at::Tensor & noise, const at::Scalar & lower, const at::Scalar & upper, bool training, bool self_is_result) {
  20864. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20865. auto maybe_layer = maybeCurrentDynamicLayer();
  20866. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20867. int64_t cur_level = maybe_layer->layerId();
  20868. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(noise, cur_level)) {
  20869. return at::_ops::rrelu_with_noise_backward::call(grad_output, self, noise, lower, upper, training, self_is_result);
  20870. }
  20871. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  20872. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20873. auto [noise_value, noise_bdim] = unwrapTensorAtLevel(noise, cur_level);
  20874. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim, noise_value, noise_bdim, lower, upper, training, self_is_result);
  20875. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20876. }
  20877. template <typename batch_rule_t, batch_rule_t batch_rule>
  20878. at::Tensor softplus_generated_plumbing(const at::Tensor & self, const at::Scalar & beta, const at::Scalar & threshold) {
  20879. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20880. auto maybe_layer = maybeCurrentDynamicLayer();
  20881. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20882. int64_t cur_level = maybe_layer->layerId();
  20883. if (!isBatchedAtLevel(self, cur_level)) {
  20884. return at::_ops::softplus::call(self, beta, threshold);
  20885. }
  20886. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20887. auto results = batch_rule(self_value, self_bdim, beta, threshold);
  20888. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20889. }
  20890. template <typename batch_rule_t, batch_rule_t batch_rule>
  20891. at::Tensor softplus_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self, const at::Scalar & beta, const at::Scalar & threshold) {
  20892. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20893. auto maybe_layer = maybeCurrentDynamicLayer();
  20894. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20895. int64_t cur_level = maybe_layer->layerId();
  20896. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level)) {
  20897. return at::_ops::softplus_backward::call(grad_output, self, beta, threshold);
  20898. }
  20899. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  20900. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20901. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim, beta, threshold);
  20902. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20903. }
  20904. template <typename batch_rule_t, batch_rule_t batch_rule>
  20905. at::Tensor softshrink_generated_plumbing(const at::Tensor & self, const at::Scalar & lambd) {
  20906. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20907. auto maybe_layer = maybeCurrentDynamicLayer();
  20908. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20909. int64_t cur_level = maybe_layer->layerId();
  20910. if (!isBatchedAtLevel(self, cur_level)) {
  20911. return at::_ops::softshrink::call(self, lambd);
  20912. }
  20913. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20914. auto results = batch_rule(self_value, self_bdim, lambd);
  20915. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20916. }
  20917. template <typename batch_rule_t, batch_rule_t batch_rule>
  20918. at::Tensor softshrink_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self, const at::Scalar & lambd) {
  20919. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20920. auto maybe_layer = maybeCurrentDynamicLayer();
  20921. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20922. int64_t cur_level = maybe_layer->layerId();
  20923. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level)) {
  20924. return at::_ops::softshrink_backward::call(grad_output, self, lambd);
  20925. }
  20926. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  20927. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20928. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim, lambd);
  20929. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20930. }
  20931. template <typename batch_rule_t, batch_rule_t batch_rule>
  20932. at::Tensor adaptive_avg_pool2d_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef output_size) {
  20933. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20934. auto maybe_layer = maybeCurrentDynamicLayer();
  20935. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20936. int64_t cur_level = maybe_layer->layerId();
  20937. if (!isBatchedAtLevel(self, cur_level)) {
  20938. return at::_ops::adaptive_avg_pool2d::call(self, output_size);
  20939. }
  20940. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20941. auto results = batch_rule(self_value, self_bdim, output_size);
  20942. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20943. }
  20944. template <typename batch_rule_t, batch_rule_t batch_rule>
  20945. at::Tensor mkldnn_adaptive_avg_pool2d_generated_plumbing(const at::Tensor & self, at::IntArrayRef output_size) {
  20946. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20947. auto maybe_layer = maybeCurrentDynamicLayer();
  20948. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20949. int64_t cur_level = maybe_layer->layerId();
  20950. if (!isBatchedAtLevel(self, cur_level)) {
  20951. return at::_ops::mkldnn_adaptive_avg_pool2d::call(self, output_size);
  20952. }
  20953. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20954. auto results = batch_rule(self_value, self_bdim, output_size);
  20955. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20956. }
  20957. template <typename batch_rule_t, batch_rule_t batch_rule>
  20958. at::Tensor mkldnn_adaptive_avg_pool2d_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self) {
  20959. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20960. auto maybe_layer = maybeCurrentDynamicLayer();
  20961. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20962. int64_t cur_level = maybe_layer->layerId();
  20963. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level)) {
  20964. return at::_ops::mkldnn_adaptive_avg_pool2d_backward::call(grad_output, self);
  20965. }
  20966. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  20967. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20968. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim);
  20969. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20970. }
  20971. template <typename batch_rule_t, batch_rule_t batch_rule>
  20972. at::Tensor _adaptive_avg_pool2d_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef output_size) {
  20973. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20974. auto maybe_layer = maybeCurrentDynamicLayer();
  20975. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20976. int64_t cur_level = maybe_layer->layerId();
  20977. if (!isBatchedAtLevel(self, cur_level)) {
  20978. return at::_ops::_adaptive_avg_pool2d::call(self, output_size);
  20979. }
  20980. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20981. auto results = batch_rule(self_value, self_bdim, output_size);
  20982. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20983. }
  20984. template <typename batch_rule_t, batch_rule_t batch_rule>
  20985. at::Tensor _adaptive_avg_pool2d_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self) {
  20986. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  20987. auto maybe_layer = maybeCurrentDynamicLayer();
  20988. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  20989. int64_t cur_level = maybe_layer->layerId();
  20990. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level)) {
  20991. return at::_ops::_adaptive_avg_pool2d_backward::call(grad_output, self);
  20992. }
  20993. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  20994. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  20995. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim);
  20996. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  20997. }
  20998. template <typename batch_rule_t, batch_rule_t batch_rule>
  20999. at::Tensor adaptive_avg_pool3d_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef output_size) {
  21000. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21001. auto maybe_layer = maybeCurrentDynamicLayer();
  21002. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21003. int64_t cur_level = maybe_layer->layerId();
  21004. if (!isBatchedAtLevel(self, cur_level)) {
  21005. return at::_ops::adaptive_avg_pool3d::call(self, output_size);
  21006. }
  21007. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21008. auto results = batch_rule(self_value, self_bdim, output_size);
  21009. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21010. }
  21011. template <typename batch_rule_t, batch_rule_t batch_rule>
  21012. at::Tensor _adaptive_avg_pool3d_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef output_size) {
  21013. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21014. auto maybe_layer = maybeCurrentDynamicLayer();
  21015. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21016. int64_t cur_level = maybe_layer->layerId();
  21017. if (!isBatchedAtLevel(self, cur_level)) {
  21018. return at::_ops::_adaptive_avg_pool3d::call(self, output_size);
  21019. }
  21020. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21021. auto results = batch_rule(self_value, self_bdim, output_size);
  21022. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21023. }
  21024. template <typename batch_rule_t, batch_rule_t batch_rule>
  21025. at::Tensor _adaptive_avg_pool3d_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self) {
  21026. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21027. auto maybe_layer = maybeCurrentDynamicLayer();
  21028. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21029. int64_t cur_level = maybe_layer->layerId();
  21030. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level)) {
  21031. return at::_ops::_adaptive_avg_pool3d_backward::call(grad_output, self);
  21032. }
  21033. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  21034. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21035. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim);
  21036. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21037. }
  21038. template <typename batch_rule_t, batch_rule_t batch_rule>
  21039. ::std::tuple<at::Tensor,at::Tensor> adaptive_max_pool2d_generated_plumbing(const at::Tensor & self, at::IntArrayRef output_size) {
  21040. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21041. auto maybe_layer = maybeCurrentDynamicLayer();
  21042. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21043. int64_t cur_level = maybe_layer->layerId();
  21044. if (!isBatchedAtLevel(self, cur_level)) {
  21045. return at::_ops::adaptive_max_pool2d::call(self, output_size);
  21046. }
  21047. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21048. auto results = batch_rule(self_value, self_bdim, output_size);
  21049. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  21050. }
  21051. template <typename batch_rule_t, batch_rule_t batch_rule>
  21052. at::Tensor adaptive_max_pool2d_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self, const at::Tensor & indices) {
  21053. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21054. auto maybe_layer = maybeCurrentDynamicLayer();
  21055. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21056. int64_t cur_level = maybe_layer->layerId();
  21057. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(indices, cur_level)) {
  21058. return at::_ops::adaptive_max_pool2d_backward::call(grad_output, self, indices);
  21059. }
  21060. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  21061. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21062. auto [indices_value, indices_bdim] = unwrapTensorAtLevel(indices, cur_level);
  21063. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim, indices_value, indices_bdim);
  21064. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21065. }
  21066. template <typename batch_rule_t, batch_rule_t batch_rule>
  21067. ::std::tuple<at::Tensor,at::Tensor> adaptive_max_pool3d_generated_plumbing(const at::Tensor & self, at::IntArrayRef output_size) {
  21068. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21069. auto maybe_layer = maybeCurrentDynamicLayer();
  21070. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21071. int64_t cur_level = maybe_layer->layerId();
  21072. if (!isBatchedAtLevel(self, cur_level)) {
  21073. return at::_ops::adaptive_max_pool3d::call(self, output_size);
  21074. }
  21075. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21076. auto results = batch_rule(self_value, self_bdim, output_size);
  21077. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  21078. }
  21079. template <typename batch_rule_t, batch_rule_t batch_rule>
  21080. at::Tensor adaptive_max_pool3d_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self, const at::Tensor & indices) {
  21081. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21082. auto maybe_layer = maybeCurrentDynamicLayer();
  21083. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21084. int64_t cur_level = maybe_layer->layerId();
  21085. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(indices, cur_level)) {
  21086. return at::_ops::adaptive_max_pool3d_backward::call(grad_output, self, indices);
  21087. }
  21088. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  21089. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21090. auto [indices_value, indices_bdim] = unwrapTensorAtLevel(indices, cur_level);
  21091. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim, indices_value, indices_bdim);
  21092. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21093. }
  21094. template <typename batch_rule_t, batch_rule_t batch_rule>
  21095. at::Tensor avg_pool2d_generated_plumbing(const at::Tensor & self, at::IntArrayRef kernel_size, at::IntArrayRef stride, at::IntArrayRef padding, bool ceil_mode, bool count_include_pad, ::std::optional<int64_t> divisor_override) {
  21096. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21097. auto maybe_layer = maybeCurrentDynamicLayer();
  21098. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21099. int64_t cur_level = maybe_layer->layerId();
  21100. if (!isBatchedAtLevel(self, cur_level)) {
  21101. return at::_ops::avg_pool2d::call(self, kernel_size, stride, padding, ceil_mode, count_include_pad, divisor_override);
  21102. }
  21103. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21104. auto results = batch_rule(self_value, self_bdim, kernel_size, stride, padding, ceil_mode, count_include_pad, divisor_override);
  21105. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21106. }
  21107. template <typename batch_rule_t, batch_rule_t batch_rule>
  21108. at::Tensor avg_pool2d_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self, at::IntArrayRef kernel_size, at::IntArrayRef stride, at::IntArrayRef padding, bool ceil_mode, bool count_include_pad, ::std::optional<int64_t> divisor_override) {
  21109. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21110. auto maybe_layer = maybeCurrentDynamicLayer();
  21111. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21112. int64_t cur_level = maybe_layer->layerId();
  21113. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level)) {
  21114. return at::_ops::avg_pool2d_backward::call(grad_output, self, kernel_size, stride, padding, ceil_mode, count_include_pad, divisor_override);
  21115. }
  21116. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  21117. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21118. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim, kernel_size, stride, padding, ceil_mode, count_include_pad, divisor_override);
  21119. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21120. }
  21121. template <typename batch_rule_t, batch_rule_t batch_rule>
  21122. at::Tensor avg_pool3d_generated_plumbing(const at::Tensor & self, at::IntArrayRef kernel_size, at::IntArrayRef stride, at::IntArrayRef padding, bool ceil_mode, bool count_include_pad, ::std::optional<int64_t> divisor_override) {
  21123. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21124. auto maybe_layer = maybeCurrentDynamicLayer();
  21125. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21126. int64_t cur_level = maybe_layer->layerId();
  21127. if (!isBatchedAtLevel(self, cur_level)) {
  21128. return at::_ops::avg_pool3d::call(self, kernel_size, stride, padding, ceil_mode, count_include_pad, divisor_override);
  21129. }
  21130. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21131. auto results = batch_rule(self_value, self_bdim, kernel_size, stride, padding, ceil_mode, count_include_pad, divisor_override);
  21132. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21133. }
  21134. template <typename batch_rule_t, batch_rule_t batch_rule>
  21135. at::Tensor avg_pool3d_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self, at::IntArrayRef kernel_size, at::IntArrayRef stride, at::IntArrayRef padding, bool ceil_mode, bool count_include_pad, ::std::optional<int64_t> divisor_override) {
  21136. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21137. auto maybe_layer = maybeCurrentDynamicLayer();
  21138. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21139. int64_t cur_level = maybe_layer->layerId();
  21140. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level)) {
  21141. return at::_ops::avg_pool3d_backward::call(grad_output, self, kernel_size, stride, padding, ceil_mode, count_include_pad, divisor_override);
  21142. }
  21143. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  21144. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21145. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim, kernel_size, stride, padding, ceil_mode, count_include_pad, divisor_override);
  21146. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21147. }
  21148. template <typename batch_rule_t, batch_rule_t batch_rule>
  21149. ::std::tuple<at::Tensor,at::Tensor> fractional_max_pool2d_generated_plumbing(const at::Tensor & self, at::IntArrayRef kernel_size, at::IntArrayRef output_size, const at::Tensor & random_samples) {
  21150. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21151. auto maybe_layer = maybeCurrentDynamicLayer();
  21152. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21153. int64_t cur_level = maybe_layer->layerId();
  21154. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(random_samples, cur_level)) {
  21155. return at::_ops::fractional_max_pool2d::call(self, kernel_size, output_size, random_samples);
  21156. }
  21157. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21158. auto [random_samples_value, random_samples_bdim] = unwrapTensorAtLevel(random_samples, cur_level);
  21159. auto results = batch_rule(self_value, self_bdim, kernel_size, output_size, random_samples_value, random_samples_bdim);
  21160. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  21161. }
  21162. template <typename batch_rule_t, batch_rule_t batch_rule>
  21163. at::Tensor fractional_max_pool2d_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self, at::IntArrayRef kernel_size, at::IntArrayRef output_size, const at::Tensor & indices) {
  21164. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21165. auto maybe_layer = maybeCurrentDynamicLayer();
  21166. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21167. int64_t cur_level = maybe_layer->layerId();
  21168. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(indices, cur_level)) {
  21169. return at::_ops::fractional_max_pool2d_backward::call(grad_output, self, kernel_size, output_size, indices);
  21170. }
  21171. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  21172. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21173. auto [indices_value, indices_bdim] = unwrapTensorAtLevel(indices, cur_level);
  21174. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim, kernel_size, output_size, indices_value, indices_bdim);
  21175. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21176. }
  21177. template <typename batch_rule_t, batch_rule_t batch_rule>
  21178. ::std::tuple<at::Tensor,at::Tensor> fractional_max_pool3d_generated_plumbing(const at::Tensor & self, at::IntArrayRef kernel_size, at::IntArrayRef output_size, const at::Tensor & random_samples) {
  21179. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21180. auto maybe_layer = maybeCurrentDynamicLayer();
  21181. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21182. int64_t cur_level = maybe_layer->layerId();
  21183. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(random_samples, cur_level)) {
  21184. return at::_ops::fractional_max_pool3d::call(self, kernel_size, output_size, random_samples);
  21185. }
  21186. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21187. auto [random_samples_value, random_samples_bdim] = unwrapTensorAtLevel(random_samples, cur_level);
  21188. auto results = batch_rule(self_value, self_bdim, kernel_size, output_size, random_samples_value, random_samples_bdim);
  21189. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  21190. }
  21191. template <typename batch_rule_t, batch_rule_t batch_rule>
  21192. at::Tensor fractional_max_pool3d_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self, at::IntArrayRef kernel_size, at::IntArrayRef output_size, const at::Tensor & indices) {
  21193. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21194. auto maybe_layer = maybeCurrentDynamicLayer();
  21195. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21196. int64_t cur_level = maybe_layer->layerId();
  21197. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(indices, cur_level)) {
  21198. return at::_ops::fractional_max_pool3d_backward::call(grad_output, self, kernel_size, output_size, indices);
  21199. }
  21200. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  21201. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21202. auto [indices_value, indices_bdim] = unwrapTensorAtLevel(indices, cur_level);
  21203. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim, kernel_size, output_size, indices_value, indices_bdim);
  21204. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21205. }
  21206. template <typename batch_rule_t, batch_rule_t batch_rule>
  21207. ::std::tuple<at::Tensor,at::Tensor> max_pool2d_with_indices_generated_plumbing(const at::Tensor & self, at::IntArrayRef kernel_size, at::IntArrayRef stride, at::IntArrayRef padding, at::IntArrayRef dilation, bool ceil_mode) {
  21208. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21209. auto maybe_layer = maybeCurrentDynamicLayer();
  21210. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21211. int64_t cur_level = maybe_layer->layerId();
  21212. if (!isBatchedAtLevel(self, cur_level)) {
  21213. return at::_ops::max_pool2d_with_indices::call(self, kernel_size, stride, padding, dilation, ceil_mode);
  21214. }
  21215. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21216. auto results = batch_rule(self_value, self_bdim, kernel_size, stride, padding, dilation, ceil_mode);
  21217. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  21218. }
  21219. template <typename batch_rule_t, batch_rule_t batch_rule>
  21220. at::Tensor max_pool2d_with_indices_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self, at::IntArrayRef kernel_size, at::IntArrayRef stride, at::IntArrayRef padding, at::IntArrayRef dilation, bool ceil_mode, const at::Tensor & indices) {
  21221. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21222. auto maybe_layer = maybeCurrentDynamicLayer();
  21223. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21224. int64_t cur_level = maybe_layer->layerId();
  21225. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(indices, cur_level)) {
  21226. return at::_ops::max_pool2d_with_indices_backward::call(grad_output, self, kernel_size, stride, padding, dilation, ceil_mode, indices);
  21227. }
  21228. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  21229. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21230. auto [indices_value, indices_bdim] = unwrapTensorAtLevel(indices, cur_level);
  21231. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim, kernel_size, stride, padding, dilation, ceil_mode, indices_value, indices_bdim);
  21232. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21233. }
  21234. template <typename batch_rule_t, batch_rule_t batch_rule>
  21235. ::std::tuple<at::Tensor,at::Tensor> max_pool3d_with_indices_generated_plumbing(const at::Tensor & self, at::IntArrayRef kernel_size, at::IntArrayRef stride, at::IntArrayRef padding, at::IntArrayRef dilation, bool ceil_mode) {
  21236. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21237. auto maybe_layer = maybeCurrentDynamicLayer();
  21238. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21239. int64_t cur_level = maybe_layer->layerId();
  21240. if (!isBatchedAtLevel(self, cur_level)) {
  21241. return at::_ops::max_pool3d_with_indices::call(self, kernel_size, stride, padding, dilation, ceil_mode);
  21242. }
  21243. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21244. auto results = batch_rule(self_value, self_bdim, kernel_size, stride, padding, dilation, ceil_mode);
  21245. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  21246. }
  21247. template <typename batch_rule_t, batch_rule_t batch_rule>
  21248. at::Tensor max_pool3d_with_indices_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self, at::IntArrayRef kernel_size, at::IntArrayRef stride, at::IntArrayRef padding, at::IntArrayRef dilation, bool ceil_mode, const at::Tensor & indices) {
  21249. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21250. auto maybe_layer = maybeCurrentDynamicLayer();
  21251. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21252. int64_t cur_level = maybe_layer->layerId();
  21253. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(indices, cur_level)) {
  21254. return at::_ops::max_pool3d_with_indices_backward::call(grad_output, self, kernel_size, stride, padding, dilation, ceil_mode, indices);
  21255. }
  21256. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  21257. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21258. auto [indices_value, indices_bdim] = unwrapTensorAtLevel(indices, cur_level);
  21259. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim, kernel_size, stride, padding, dilation, ceil_mode, indices_value, indices_bdim);
  21260. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21261. }
  21262. template <typename batch_rule_t, batch_rule_t batch_rule>
  21263. at::Tensor max_unpool2d_generated_plumbing(const at::Tensor & self, const at::Tensor & indices, c10::SymIntArrayRef output_size) {
  21264. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21265. auto maybe_layer = maybeCurrentDynamicLayer();
  21266. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21267. int64_t cur_level = maybe_layer->layerId();
  21268. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(indices, cur_level)) {
  21269. return at::_ops::max_unpool2d::call(self, indices, output_size);
  21270. }
  21271. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21272. auto [indices_value, indices_bdim] = unwrapTensorAtLevel(indices, cur_level);
  21273. auto results = batch_rule(self_value, self_bdim, indices_value, indices_bdim, output_size);
  21274. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21275. }
  21276. template <typename batch_rule_t, batch_rule_t batch_rule>
  21277. at::Tensor max_unpool3d_generated_plumbing(const at::Tensor & self, const at::Tensor & indices, c10::SymIntArrayRef output_size, at::IntArrayRef stride, at::IntArrayRef padding) {
  21278. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21279. auto maybe_layer = maybeCurrentDynamicLayer();
  21280. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21281. int64_t cur_level = maybe_layer->layerId();
  21282. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(indices, cur_level)) {
  21283. return at::_ops::max_unpool3d::call(self, indices, output_size, stride, padding);
  21284. }
  21285. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21286. auto [indices_value, indices_bdim] = unwrapTensorAtLevel(indices, cur_level);
  21287. auto results = batch_rule(self_value, self_bdim, indices_value, indices_bdim, output_size, stride, padding);
  21288. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21289. }
  21290. template <typename batch_rule_t, batch_rule_t batch_rule>
  21291. at::Tensor reflection_pad1d_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef padding) {
  21292. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21293. auto maybe_layer = maybeCurrentDynamicLayer();
  21294. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21295. int64_t cur_level = maybe_layer->layerId();
  21296. if (!isBatchedAtLevel(self, cur_level)) {
  21297. return at::_ops::reflection_pad1d::call(self, padding);
  21298. }
  21299. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21300. auto results = batch_rule(self_value, self_bdim, padding);
  21301. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21302. }
  21303. template <typename batch_rule_t, batch_rule_t batch_rule>
  21304. at::Tensor reflection_pad1d_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self, c10::SymIntArrayRef padding) {
  21305. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21306. auto maybe_layer = maybeCurrentDynamicLayer();
  21307. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21308. int64_t cur_level = maybe_layer->layerId();
  21309. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level)) {
  21310. return at::_ops::reflection_pad1d_backward::call(grad_output, self, padding);
  21311. }
  21312. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  21313. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21314. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim, padding);
  21315. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21316. }
  21317. template <typename batch_rule_t, batch_rule_t batch_rule>
  21318. at::Tensor reflection_pad2d_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef padding) {
  21319. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21320. auto maybe_layer = maybeCurrentDynamicLayer();
  21321. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21322. int64_t cur_level = maybe_layer->layerId();
  21323. if (!isBatchedAtLevel(self, cur_level)) {
  21324. return at::_ops::reflection_pad2d::call(self, padding);
  21325. }
  21326. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21327. auto results = batch_rule(self_value, self_bdim, padding);
  21328. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21329. }
  21330. template <typename batch_rule_t, batch_rule_t batch_rule>
  21331. at::Tensor reflection_pad2d_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self, c10::SymIntArrayRef padding) {
  21332. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21333. auto maybe_layer = maybeCurrentDynamicLayer();
  21334. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21335. int64_t cur_level = maybe_layer->layerId();
  21336. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level)) {
  21337. return at::_ops::reflection_pad2d_backward::call(grad_output, self, padding);
  21338. }
  21339. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  21340. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21341. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim, padding);
  21342. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21343. }
  21344. template <typename batch_rule_t, batch_rule_t batch_rule>
  21345. at::Tensor reflection_pad3d_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef padding) {
  21346. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21347. auto maybe_layer = maybeCurrentDynamicLayer();
  21348. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21349. int64_t cur_level = maybe_layer->layerId();
  21350. if (!isBatchedAtLevel(self, cur_level)) {
  21351. return at::_ops::reflection_pad3d::call(self, padding);
  21352. }
  21353. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21354. auto results = batch_rule(self_value, self_bdim, padding);
  21355. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21356. }
  21357. template <typename batch_rule_t, batch_rule_t batch_rule>
  21358. at::Tensor reflection_pad3d_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self, c10::SymIntArrayRef padding) {
  21359. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21360. auto maybe_layer = maybeCurrentDynamicLayer();
  21361. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21362. int64_t cur_level = maybe_layer->layerId();
  21363. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level)) {
  21364. return at::_ops::reflection_pad3d_backward::call(grad_output, self, padding);
  21365. }
  21366. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  21367. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21368. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim, padding);
  21369. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21370. }
  21371. template <typename batch_rule_t, batch_rule_t batch_rule>
  21372. at::Tensor replication_pad1d_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef padding) {
  21373. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21374. auto maybe_layer = maybeCurrentDynamicLayer();
  21375. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21376. int64_t cur_level = maybe_layer->layerId();
  21377. if (!isBatchedAtLevel(self, cur_level)) {
  21378. return at::_ops::replication_pad1d::call(self, padding);
  21379. }
  21380. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21381. auto results = batch_rule(self_value, self_bdim, padding);
  21382. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21383. }
  21384. template <typename batch_rule_t, batch_rule_t batch_rule>
  21385. at::Tensor replication_pad1d_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self, c10::SymIntArrayRef padding) {
  21386. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21387. auto maybe_layer = maybeCurrentDynamicLayer();
  21388. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21389. int64_t cur_level = maybe_layer->layerId();
  21390. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level)) {
  21391. return at::_ops::replication_pad1d_backward::call(grad_output, self, padding);
  21392. }
  21393. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  21394. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21395. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim, padding);
  21396. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21397. }
  21398. template <typename batch_rule_t, batch_rule_t batch_rule>
  21399. at::Tensor replication_pad2d_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef padding) {
  21400. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21401. auto maybe_layer = maybeCurrentDynamicLayer();
  21402. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21403. int64_t cur_level = maybe_layer->layerId();
  21404. if (!isBatchedAtLevel(self, cur_level)) {
  21405. return at::_ops::replication_pad2d::call(self, padding);
  21406. }
  21407. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21408. auto results = batch_rule(self_value, self_bdim, padding);
  21409. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21410. }
  21411. template <typename batch_rule_t, batch_rule_t batch_rule>
  21412. at::Tensor replication_pad2d_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self, c10::SymIntArrayRef padding) {
  21413. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21414. auto maybe_layer = maybeCurrentDynamicLayer();
  21415. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21416. int64_t cur_level = maybe_layer->layerId();
  21417. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level)) {
  21418. return at::_ops::replication_pad2d_backward::call(grad_output, self, padding);
  21419. }
  21420. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  21421. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21422. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim, padding);
  21423. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21424. }
  21425. template <typename batch_rule_t, batch_rule_t batch_rule>
  21426. at::Tensor replication_pad3d_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef padding) {
  21427. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21428. auto maybe_layer = maybeCurrentDynamicLayer();
  21429. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21430. int64_t cur_level = maybe_layer->layerId();
  21431. if (!isBatchedAtLevel(self, cur_level)) {
  21432. return at::_ops::replication_pad3d::call(self, padding);
  21433. }
  21434. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21435. auto results = batch_rule(self_value, self_bdim, padding);
  21436. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21437. }
  21438. template <typename batch_rule_t, batch_rule_t batch_rule>
  21439. at::Tensor replication_pad3d_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self, c10::SymIntArrayRef padding) {
  21440. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21441. auto maybe_layer = maybeCurrentDynamicLayer();
  21442. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21443. int64_t cur_level = maybe_layer->layerId();
  21444. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level)) {
  21445. return at::_ops::replication_pad3d_backward::call(grad_output, self, padding);
  21446. }
  21447. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  21448. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21449. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim, padding);
  21450. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21451. }
  21452. template <typename batch_rule_t, batch_rule_t batch_rule>
  21453. at::Tensor _pad_circular_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef pad) {
  21454. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21455. auto maybe_layer = maybeCurrentDynamicLayer();
  21456. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21457. int64_t cur_level = maybe_layer->layerId();
  21458. if (!isBatchedAtLevel(self, cur_level)) {
  21459. return at::_ops::_pad_circular::call(self, pad);
  21460. }
  21461. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21462. auto results = batch_rule(self_value, self_bdim, pad);
  21463. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21464. }
  21465. template <typename batch_rule_t, batch_rule_t batch_rule>
  21466. at::Tensor _pad_enum_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef pad, int64_t mode, ::std::optional<double> value) {
  21467. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21468. auto maybe_layer = maybeCurrentDynamicLayer();
  21469. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21470. int64_t cur_level = maybe_layer->layerId();
  21471. if (!isBatchedAtLevel(self, cur_level)) {
  21472. return at::_ops::_pad_enum::call(self, pad, mode, value);
  21473. }
  21474. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21475. auto results = batch_rule(self_value, self_bdim, pad, mode, value);
  21476. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21477. }
  21478. template <typename batch_rule_t, batch_rule_t batch_rule>
  21479. at::Tensor pad_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef pad, c10::string_view mode, ::std::optional<double> value) {
  21480. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21481. auto maybe_layer = maybeCurrentDynamicLayer();
  21482. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21483. int64_t cur_level = maybe_layer->layerId();
  21484. if (!isBatchedAtLevel(self, cur_level)) {
  21485. return at::_ops::pad::call(self, pad, mode, value);
  21486. }
  21487. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21488. auto results = batch_rule(self_value, self_bdim, pad, mode, value);
  21489. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21490. }
  21491. template <typename batch_rule_t, batch_rule_t batch_rule>
  21492. at::Tensor upsample_linear1d_vec_generated_plumbing(const at::Tensor & input, at::OptionalSymIntArrayRef output_size, bool align_corners, ::std::optional<at::ArrayRef<double>> scale_factors) {
  21493. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21494. auto maybe_layer = maybeCurrentDynamicLayer();
  21495. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21496. int64_t cur_level = maybe_layer->layerId();
  21497. if (!isBatchedAtLevel(input, cur_level)) {
  21498. return at::_ops::upsample_linear1d_vec::call(input, output_size, align_corners, scale_factors);
  21499. }
  21500. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  21501. auto results = batch_rule(input_value, input_bdim, output_size, align_corners, scale_factors);
  21502. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21503. }
  21504. template <typename batch_rule_t, batch_rule_t batch_rule>
  21505. at::Tensor upsample_bilinear2d_vec_generated_plumbing(const at::Tensor & input, at::OptionalSymIntArrayRef output_size, bool align_corners, ::std::optional<at::ArrayRef<double>> scale_factors) {
  21506. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21507. auto maybe_layer = maybeCurrentDynamicLayer();
  21508. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21509. int64_t cur_level = maybe_layer->layerId();
  21510. if (!isBatchedAtLevel(input, cur_level)) {
  21511. return at::_ops::upsample_bilinear2d_vec::call(input, output_size, align_corners, scale_factors);
  21512. }
  21513. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  21514. auto results = batch_rule(input_value, input_bdim, output_size, align_corners, scale_factors);
  21515. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21516. }
  21517. template <typename batch_rule_t, batch_rule_t batch_rule>
  21518. at::Tensor _upsample_bilinear2d_aa_vec_generated_plumbing(const at::Tensor & input, at::OptionalSymIntArrayRef output_size, bool align_corners, ::std::optional<at::ArrayRef<double>> scale_factors) {
  21519. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21520. auto maybe_layer = maybeCurrentDynamicLayer();
  21521. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21522. int64_t cur_level = maybe_layer->layerId();
  21523. if (!isBatchedAtLevel(input, cur_level)) {
  21524. return at::_ops::_upsample_bilinear2d_aa_vec::call(input, output_size, align_corners, scale_factors);
  21525. }
  21526. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  21527. auto results = batch_rule(input_value, input_bdim, output_size, align_corners, scale_factors);
  21528. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21529. }
  21530. template <typename batch_rule_t, batch_rule_t batch_rule>
  21531. at::Tensor upsample_trilinear3d_vec_generated_plumbing(const at::Tensor & input, at::OptionalSymIntArrayRef output_size, bool align_corners, ::std::optional<at::ArrayRef<double>> scale_factors) {
  21532. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21533. auto maybe_layer = maybeCurrentDynamicLayer();
  21534. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21535. int64_t cur_level = maybe_layer->layerId();
  21536. if (!isBatchedAtLevel(input, cur_level)) {
  21537. return at::_ops::upsample_trilinear3d_vec::call(input, output_size, align_corners, scale_factors);
  21538. }
  21539. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  21540. auto results = batch_rule(input_value, input_bdim, output_size, align_corners, scale_factors);
  21541. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21542. }
  21543. template <typename batch_rule_t, batch_rule_t batch_rule>
  21544. at::Tensor upsample_bicubic2d_vec_generated_plumbing(const at::Tensor & input, at::OptionalSymIntArrayRef output_size, bool align_corners, ::std::optional<at::ArrayRef<double>> scale_factors) {
  21545. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21546. auto maybe_layer = maybeCurrentDynamicLayer();
  21547. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21548. int64_t cur_level = maybe_layer->layerId();
  21549. if (!isBatchedAtLevel(input, cur_level)) {
  21550. return at::_ops::upsample_bicubic2d_vec::call(input, output_size, align_corners, scale_factors);
  21551. }
  21552. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  21553. auto results = batch_rule(input_value, input_bdim, output_size, align_corners, scale_factors);
  21554. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21555. }
  21556. template <typename batch_rule_t, batch_rule_t batch_rule>
  21557. at::Tensor _upsample_bicubic2d_aa_vec_generated_plumbing(const at::Tensor & input, at::OptionalSymIntArrayRef output_size, bool align_corners, ::std::optional<at::ArrayRef<double>> scale_factors) {
  21558. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21559. auto maybe_layer = maybeCurrentDynamicLayer();
  21560. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21561. int64_t cur_level = maybe_layer->layerId();
  21562. if (!isBatchedAtLevel(input, cur_level)) {
  21563. return at::_ops::_upsample_bicubic2d_aa_vec::call(input, output_size, align_corners, scale_factors);
  21564. }
  21565. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  21566. auto results = batch_rule(input_value, input_bdim, output_size, align_corners, scale_factors);
  21567. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21568. }
  21569. template <typename batch_rule_t, batch_rule_t batch_rule>
  21570. at::Tensor upsample_nearest1d_vec_generated_plumbing(const at::Tensor & input, at::OptionalSymIntArrayRef output_size, ::std::optional<at::ArrayRef<double>> scale_factors) {
  21571. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21572. auto maybe_layer = maybeCurrentDynamicLayer();
  21573. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21574. int64_t cur_level = maybe_layer->layerId();
  21575. if (!isBatchedAtLevel(input, cur_level)) {
  21576. return at::_ops::upsample_nearest1d_vec::call(input, output_size, scale_factors);
  21577. }
  21578. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  21579. auto results = batch_rule(input_value, input_bdim, output_size, scale_factors);
  21580. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21581. }
  21582. template <typename batch_rule_t, batch_rule_t batch_rule>
  21583. at::Tensor _upsample_nearest_exact1d_vec_generated_plumbing(const at::Tensor & input, at::OptionalSymIntArrayRef output_size, ::std::optional<at::ArrayRef<double>> scale_factors) {
  21584. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21585. auto maybe_layer = maybeCurrentDynamicLayer();
  21586. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21587. int64_t cur_level = maybe_layer->layerId();
  21588. if (!isBatchedAtLevel(input, cur_level)) {
  21589. return at::_ops::_upsample_nearest_exact1d_vec::call(input, output_size, scale_factors);
  21590. }
  21591. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  21592. auto results = batch_rule(input_value, input_bdim, output_size, scale_factors);
  21593. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21594. }
  21595. template <typename batch_rule_t, batch_rule_t batch_rule>
  21596. at::Tensor upsample_nearest2d_vec_generated_plumbing(const at::Tensor & input, at::OptionalSymIntArrayRef output_size, ::std::optional<at::ArrayRef<double>> scale_factors) {
  21597. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21598. auto maybe_layer = maybeCurrentDynamicLayer();
  21599. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21600. int64_t cur_level = maybe_layer->layerId();
  21601. if (!isBatchedAtLevel(input, cur_level)) {
  21602. return at::_ops::upsample_nearest2d_vec::call(input, output_size, scale_factors);
  21603. }
  21604. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  21605. auto results = batch_rule(input_value, input_bdim, output_size, scale_factors);
  21606. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21607. }
  21608. template <typename batch_rule_t, batch_rule_t batch_rule>
  21609. at::Tensor _upsample_nearest_exact2d_vec_generated_plumbing(const at::Tensor & input, at::OptionalSymIntArrayRef output_size, ::std::optional<at::ArrayRef<double>> scale_factors) {
  21610. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21611. auto maybe_layer = maybeCurrentDynamicLayer();
  21612. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21613. int64_t cur_level = maybe_layer->layerId();
  21614. if (!isBatchedAtLevel(input, cur_level)) {
  21615. return at::_ops::_upsample_nearest_exact2d_vec::call(input, output_size, scale_factors);
  21616. }
  21617. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  21618. auto results = batch_rule(input_value, input_bdim, output_size, scale_factors);
  21619. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21620. }
  21621. template <typename batch_rule_t, batch_rule_t batch_rule>
  21622. at::Tensor upsample_nearest3d_vec_generated_plumbing(const at::Tensor & input, at::OptionalSymIntArrayRef output_size, ::std::optional<at::ArrayRef<double>> scale_factors) {
  21623. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21624. auto maybe_layer = maybeCurrentDynamicLayer();
  21625. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21626. int64_t cur_level = maybe_layer->layerId();
  21627. if (!isBatchedAtLevel(input, cur_level)) {
  21628. return at::_ops::upsample_nearest3d_vec::call(input, output_size, scale_factors);
  21629. }
  21630. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  21631. auto results = batch_rule(input_value, input_bdim, output_size, scale_factors);
  21632. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21633. }
  21634. template <typename batch_rule_t, batch_rule_t batch_rule>
  21635. at::Tensor _upsample_nearest_exact3d_vec_generated_plumbing(const at::Tensor & input, at::OptionalSymIntArrayRef output_size, ::std::optional<at::ArrayRef<double>> scale_factors) {
  21636. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21637. auto maybe_layer = maybeCurrentDynamicLayer();
  21638. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21639. int64_t cur_level = maybe_layer->layerId();
  21640. if (!isBatchedAtLevel(input, cur_level)) {
  21641. return at::_ops::_upsample_nearest_exact3d_vec::call(input, output_size, scale_factors);
  21642. }
  21643. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  21644. auto results = batch_rule(input_value, input_bdim, output_size, scale_factors);
  21645. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21646. }
  21647. template <typename batch_rule_t, batch_rule_t batch_rule>
  21648. at::Tensor upsample_linear1d_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef output_size, bool align_corners, ::std::optional<double> scales) {
  21649. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21650. auto maybe_layer = maybeCurrentDynamicLayer();
  21651. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21652. int64_t cur_level = maybe_layer->layerId();
  21653. if (!isBatchedAtLevel(self, cur_level)) {
  21654. return at::_ops::upsample_linear1d::call(self, output_size, align_corners, scales);
  21655. }
  21656. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21657. auto results = batch_rule(self_value, self_bdim, output_size, align_corners, scales);
  21658. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21659. }
  21660. template <typename batch_rule_t, batch_rule_t batch_rule>
  21661. at::Tensor upsample_linear1d_backward_generated_plumbing(const at::Tensor & grad_output, c10::SymIntArrayRef output_size, c10::SymIntArrayRef input_size, bool align_corners, ::std::optional<double> scales) {
  21662. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21663. auto maybe_layer = maybeCurrentDynamicLayer();
  21664. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21665. int64_t cur_level = maybe_layer->layerId();
  21666. if (!isBatchedAtLevel(grad_output, cur_level)) {
  21667. return at::_ops::upsample_linear1d_backward::call(grad_output, output_size, input_size, align_corners, scales);
  21668. }
  21669. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  21670. auto results = batch_rule(grad_output_value, grad_output_bdim, output_size, input_size, align_corners, scales);
  21671. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21672. }
  21673. template <typename batch_rule_t, batch_rule_t batch_rule>
  21674. at::Tensor upsample_bilinear2d_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef output_size, bool align_corners, ::std::optional<double> scales_h, ::std::optional<double> scales_w) {
  21675. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21676. auto maybe_layer = maybeCurrentDynamicLayer();
  21677. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21678. int64_t cur_level = maybe_layer->layerId();
  21679. if (!isBatchedAtLevel(self, cur_level)) {
  21680. return at::_ops::upsample_bilinear2d::call(self, output_size, align_corners, scales_h, scales_w);
  21681. }
  21682. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21683. auto results = batch_rule(self_value, self_bdim, output_size, align_corners, scales_h, scales_w);
  21684. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21685. }
  21686. template <typename batch_rule_t, batch_rule_t batch_rule>
  21687. at::Tensor upsample_bilinear2d_backward_generated_plumbing(const at::Tensor & grad_output, c10::SymIntArrayRef output_size, c10::SymIntArrayRef input_size, bool align_corners, ::std::optional<double> scales_h, ::std::optional<double> scales_w) {
  21688. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21689. auto maybe_layer = maybeCurrentDynamicLayer();
  21690. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21691. int64_t cur_level = maybe_layer->layerId();
  21692. if (!isBatchedAtLevel(grad_output, cur_level)) {
  21693. return at::_ops::upsample_bilinear2d_backward::call(grad_output, output_size, input_size, align_corners, scales_h, scales_w);
  21694. }
  21695. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  21696. auto results = batch_rule(grad_output_value, grad_output_bdim, output_size, input_size, align_corners, scales_h, scales_w);
  21697. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21698. }
  21699. template <typename batch_rule_t, batch_rule_t batch_rule>
  21700. at::Tensor _upsample_bilinear2d_aa_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef output_size, bool align_corners, ::std::optional<double> scales_h, ::std::optional<double> scales_w) {
  21701. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21702. auto maybe_layer = maybeCurrentDynamicLayer();
  21703. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21704. int64_t cur_level = maybe_layer->layerId();
  21705. if (!isBatchedAtLevel(self, cur_level)) {
  21706. return at::_ops::_upsample_bilinear2d_aa::call(self, output_size, align_corners, scales_h, scales_w);
  21707. }
  21708. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21709. auto results = batch_rule(self_value, self_bdim, output_size, align_corners, scales_h, scales_w);
  21710. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21711. }
  21712. template <typename batch_rule_t, batch_rule_t batch_rule>
  21713. at::Tensor _upsample_bilinear2d_aa_backward_generated_plumbing(const at::Tensor & grad_output, c10::SymIntArrayRef output_size, c10::SymIntArrayRef input_size, bool align_corners, ::std::optional<double> scales_h, ::std::optional<double> scales_w) {
  21714. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21715. auto maybe_layer = maybeCurrentDynamicLayer();
  21716. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21717. int64_t cur_level = maybe_layer->layerId();
  21718. if (!isBatchedAtLevel(grad_output, cur_level)) {
  21719. return at::_ops::_upsample_bilinear2d_aa_backward::call(grad_output, output_size, input_size, align_corners, scales_h, scales_w);
  21720. }
  21721. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  21722. auto results = batch_rule(grad_output_value, grad_output_bdim, output_size, input_size, align_corners, scales_h, scales_w);
  21723. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21724. }
  21725. template <typename batch_rule_t, batch_rule_t batch_rule>
  21726. at::Tensor upsample_bicubic2d_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef output_size, bool align_corners, ::std::optional<double> scales_h, ::std::optional<double> scales_w) {
  21727. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21728. auto maybe_layer = maybeCurrentDynamicLayer();
  21729. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21730. int64_t cur_level = maybe_layer->layerId();
  21731. if (!isBatchedAtLevel(self, cur_level)) {
  21732. return at::_ops::upsample_bicubic2d::call(self, output_size, align_corners, scales_h, scales_w);
  21733. }
  21734. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21735. auto results = batch_rule(self_value, self_bdim, output_size, align_corners, scales_h, scales_w);
  21736. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21737. }
  21738. template <typename batch_rule_t, batch_rule_t batch_rule>
  21739. at::Tensor upsample_bicubic2d_backward_generated_plumbing(const at::Tensor & grad_output, c10::SymIntArrayRef output_size, c10::SymIntArrayRef input_size, bool align_corners, ::std::optional<double> scales_h, ::std::optional<double> scales_w) {
  21740. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21741. auto maybe_layer = maybeCurrentDynamicLayer();
  21742. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21743. int64_t cur_level = maybe_layer->layerId();
  21744. if (!isBatchedAtLevel(grad_output, cur_level)) {
  21745. return at::_ops::upsample_bicubic2d_backward::call(grad_output, output_size, input_size, align_corners, scales_h, scales_w);
  21746. }
  21747. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  21748. auto results = batch_rule(grad_output_value, grad_output_bdim, output_size, input_size, align_corners, scales_h, scales_w);
  21749. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21750. }
  21751. template <typename batch_rule_t, batch_rule_t batch_rule>
  21752. at::Tensor _upsample_bicubic2d_aa_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef output_size, bool align_corners, ::std::optional<double> scales_h, ::std::optional<double> scales_w) {
  21753. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21754. auto maybe_layer = maybeCurrentDynamicLayer();
  21755. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21756. int64_t cur_level = maybe_layer->layerId();
  21757. if (!isBatchedAtLevel(self, cur_level)) {
  21758. return at::_ops::_upsample_bicubic2d_aa::call(self, output_size, align_corners, scales_h, scales_w);
  21759. }
  21760. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21761. auto results = batch_rule(self_value, self_bdim, output_size, align_corners, scales_h, scales_w);
  21762. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21763. }
  21764. template <typename batch_rule_t, batch_rule_t batch_rule>
  21765. at::Tensor _upsample_bicubic2d_aa_backward_generated_plumbing(const at::Tensor & grad_output, c10::SymIntArrayRef output_size, c10::SymIntArrayRef input_size, bool align_corners, ::std::optional<double> scales_h, ::std::optional<double> scales_w) {
  21766. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21767. auto maybe_layer = maybeCurrentDynamicLayer();
  21768. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21769. int64_t cur_level = maybe_layer->layerId();
  21770. if (!isBatchedAtLevel(grad_output, cur_level)) {
  21771. return at::_ops::_upsample_bicubic2d_aa_backward::call(grad_output, output_size, input_size, align_corners, scales_h, scales_w);
  21772. }
  21773. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  21774. auto results = batch_rule(grad_output_value, grad_output_bdim, output_size, input_size, align_corners, scales_h, scales_w);
  21775. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21776. }
  21777. template <typename batch_rule_t, batch_rule_t batch_rule>
  21778. at::Tensor upsample_trilinear3d_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef output_size, bool align_corners, ::std::optional<double> scales_d, ::std::optional<double> scales_h, ::std::optional<double> scales_w) {
  21779. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21780. auto maybe_layer = maybeCurrentDynamicLayer();
  21781. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21782. int64_t cur_level = maybe_layer->layerId();
  21783. if (!isBatchedAtLevel(self, cur_level)) {
  21784. return at::_ops::upsample_trilinear3d::call(self, output_size, align_corners, scales_d, scales_h, scales_w);
  21785. }
  21786. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21787. auto results = batch_rule(self_value, self_bdim, output_size, align_corners, scales_d, scales_h, scales_w);
  21788. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21789. }
  21790. template <typename batch_rule_t, batch_rule_t batch_rule>
  21791. at::Tensor upsample_trilinear3d_backward_generated_plumbing(const at::Tensor & grad_output, c10::SymIntArrayRef output_size, c10::SymIntArrayRef input_size, bool align_corners, ::std::optional<double> scales_d, ::std::optional<double> scales_h, ::std::optional<double> scales_w) {
  21792. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21793. auto maybe_layer = maybeCurrentDynamicLayer();
  21794. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21795. int64_t cur_level = maybe_layer->layerId();
  21796. if (!isBatchedAtLevel(grad_output, cur_level)) {
  21797. return at::_ops::upsample_trilinear3d_backward::call(grad_output, output_size, input_size, align_corners, scales_d, scales_h, scales_w);
  21798. }
  21799. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  21800. auto results = batch_rule(grad_output_value, grad_output_bdim, output_size, input_size, align_corners, scales_d, scales_h, scales_w);
  21801. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21802. }
  21803. template <typename batch_rule_t, batch_rule_t batch_rule>
  21804. at::Tensor upsample_nearest1d_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef output_size, ::std::optional<double> scales) {
  21805. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21806. auto maybe_layer = maybeCurrentDynamicLayer();
  21807. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21808. int64_t cur_level = maybe_layer->layerId();
  21809. if (!isBatchedAtLevel(self, cur_level)) {
  21810. return at::_ops::upsample_nearest1d::call(self, output_size, scales);
  21811. }
  21812. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21813. auto results = batch_rule(self_value, self_bdim, output_size, scales);
  21814. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21815. }
  21816. template <typename batch_rule_t, batch_rule_t batch_rule>
  21817. at::Tensor _upsample_nearest_exact1d_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef output_size, ::std::optional<double> scales) {
  21818. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21819. auto maybe_layer = maybeCurrentDynamicLayer();
  21820. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21821. int64_t cur_level = maybe_layer->layerId();
  21822. if (!isBatchedAtLevel(self, cur_level)) {
  21823. return at::_ops::_upsample_nearest_exact1d::call(self, output_size, scales);
  21824. }
  21825. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21826. auto results = batch_rule(self_value, self_bdim, output_size, scales);
  21827. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21828. }
  21829. template <typename batch_rule_t, batch_rule_t batch_rule>
  21830. at::Tensor upsample_nearest1d_backward_generated_plumbing(const at::Tensor & grad_output, c10::SymIntArrayRef output_size, c10::SymIntArrayRef input_size, ::std::optional<double> scales) {
  21831. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21832. auto maybe_layer = maybeCurrentDynamicLayer();
  21833. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21834. int64_t cur_level = maybe_layer->layerId();
  21835. if (!isBatchedAtLevel(grad_output, cur_level)) {
  21836. return at::_ops::upsample_nearest1d_backward::call(grad_output, output_size, input_size, scales);
  21837. }
  21838. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  21839. auto results = batch_rule(grad_output_value, grad_output_bdim, output_size, input_size, scales);
  21840. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21841. }
  21842. template <typename batch_rule_t, batch_rule_t batch_rule>
  21843. at::Tensor _upsample_nearest_exact1d_backward_generated_plumbing(const at::Tensor & grad_output, c10::SymIntArrayRef output_size, c10::SymIntArrayRef input_size, ::std::optional<double> scales) {
  21844. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21845. auto maybe_layer = maybeCurrentDynamicLayer();
  21846. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21847. int64_t cur_level = maybe_layer->layerId();
  21848. if (!isBatchedAtLevel(grad_output, cur_level)) {
  21849. return at::_ops::_upsample_nearest_exact1d_backward::call(grad_output, output_size, input_size, scales);
  21850. }
  21851. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  21852. auto results = batch_rule(grad_output_value, grad_output_bdim, output_size, input_size, scales);
  21853. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21854. }
  21855. template <typename batch_rule_t, batch_rule_t batch_rule>
  21856. at::Tensor upsample_nearest2d_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef output_size, ::std::optional<double> scales_h, ::std::optional<double> scales_w) {
  21857. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21858. auto maybe_layer = maybeCurrentDynamicLayer();
  21859. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21860. int64_t cur_level = maybe_layer->layerId();
  21861. if (!isBatchedAtLevel(self, cur_level)) {
  21862. return at::_ops::upsample_nearest2d::call(self, output_size, scales_h, scales_w);
  21863. }
  21864. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21865. auto results = batch_rule(self_value, self_bdim, output_size, scales_h, scales_w);
  21866. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21867. }
  21868. template <typename batch_rule_t, batch_rule_t batch_rule>
  21869. at::Tensor _upsample_nearest_exact2d_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef output_size, ::std::optional<double> scales_h, ::std::optional<double> scales_w) {
  21870. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21871. auto maybe_layer = maybeCurrentDynamicLayer();
  21872. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21873. int64_t cur_level = maybe_layer->layerId();
  21874. if (!isBatchedAtLevel(self, cur_level)) {
  21875. return at::_ops::_upsample_nearest_exact2d::call(self, output_size, scales_h, scales_w);
  21876. }
  21877. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21878. auto results = batch_rule(self_value, self_bdim, output_size, scales_h, scales_w);
  21879. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21880. }
  21881. template <typename batch_rule_t, batch_rule_t batch_rule>
  21882. at::Tensor upsample_nearest2d_backward_generated_plumbing(const at::Tensor & grad_output, c10::SymIntArrayRef output_size, c10::SymIntArrayRef input_size, ::std::optional<double> scales_h, ::std::optional<double> scales_w) {
  21883. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21884. auto maybe_layer = maybeCurrentDynamicLayer();
  21885. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21886. int64_t cur_level = maybe_layer->layerId();
  21887. if (!isBatchedAtLevel(grad_output, cur_level)) {
  21888. return at::_ops::upsample_nearest2d_backward::call(grad_output, output_size, input_size, scales_h, scales_w);
  21889. }
  21890. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  21891. auto results = batch_rule(grad_output_value, grad_output_bdim, output_size, input_size, scales_h, scales_w);
  21892. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21893. }
  21894. template <typename batch_rule_t, batch_rule_t batch_rule>
  21895. at::Tensor _upsample_nearest_exact2d_backward_generated_plumbing(const at::Tensor & grad_output, c10::SymIntArrayRef output_size, c10::SymIntArrayRef input_size, ::std::optional<double> scales_h, ::std::optional<double> scales_w) {
  21896. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21897. auto maybe_layer = maybeCurrentDynamicLayer();
  21898. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21899. int64_t cur_level = maybe_layer->layerId();
  21900. if (!isBatchedAtLevel(grad_output, cur_level)) {
  21901. return at::_ops::_upsample_nearest_exact2d_backward::call(grad_output, output_size, input_size, scales_h, scales_w);
  21902. }
  21903. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  21904. auto results = batch_rule(grad_output_value, grad_output_bdim, output_size, input_size, scales_h, scales_w);
  21905. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21906. }
  21907. template <typename batch_rule_t, batch_rule_t batch_rule>
  21908. at::Tensor upsample_nearest3d_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef output_size, ::std::optional<double> scales_d, ::std::optional<double> scales_h, ::std::optional<double> scales_w) {
  21909. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21910. auto maybe_layer = maybeCurrentDynamicLayer();
  21911. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21912. int64_t cur_level = maybe_layer->layerId();
  21913. if (!isBatchedAtLevel(self, cur_level)) {
  21914. return at::_ops::upsample_nearest3d::call(self, output_size, scales_d, scales_h, scales_w);
  21915. }
  21916. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21917. auto results = batch_rule(self_value, self_bdim, output_size, scales_d, scales_h, scales_w);
  21918. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21919. }
  21920. template <typename batch_rule_t, batch_rule_t batch_rule>
  21921. at::Tensor _upsample_nearest_exact3d_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef output_size, ::std::optional<double> scales_d, ::std::optional<double> scales_h, ::std::optional<double> scales_w) {
  21922. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21923. auto maybe_layer = maybeCurrentDynamicLayer();
  21924. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21925. int64_t cur_level = maybe_layer->layerId();
  21926. if (!isBatchedAtLevel(self, cur_level)) {
  21927. return at::_ops::_upsample_nearest_exact3d::call(self, output_size, scales_d, scales_h, scales_w);
  21928. }
  21929. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21930. auto results = batch_rule(self_value, self_bdim, output_size, scales_d, scales_h, scales_w);
  21931. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21932. }
  21933. template <typename batch_rule_t, batch_rule_t batch_rule>
  21934. at::Tensor upsample_nearest3d_backward_generated_plumbing(const at::Tensor & grad_output, c10::SymIntArrayRef output_size, c10::SymIntArrayRef input_size, ::std::optional<double> scales_d, ::std::optional<double> scales_h, ::std::optional<double> scales_w) {
  21935. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21936. auto maybe_layer = maybeCurrentDynamicLayer();
  21937. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21938. int64_t cur_level = maybe_layer->layerId();
  21939. if (!isBatchedAtLevel(grad_output, cur_level)) {
  21940. return at::_ops::upsample_nearest3d_backward::call(grad_output, output_size, input_size, scales_d, scales_h, scales_w);
  21941. }
  21942. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  21943. auto results = batch_rule(grad_output_value, grad_output_bdim, output_size, input_size, scales_d, scales_h, scales_w);
  21944. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21945. }
  21946. template <typename batch_rule_t, batch_rule_t batch_rule>
  21947. at::Tensor _upsample_nearest_exact3d_backward_generated_plumbing(const at::Tensor & grad_output, c10::SymIntArrayRef output_size, c10::SymIntArrayRef input_size, ::std::optional<double> scales_d, ::std::optional<double> scales_h, ::std::optional<double> scales_w) {
  21948. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21949. auto maybe_layer = maybeCurrentDynamicLayer();
  21950. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21951. int64_t cur_level = maybe_layer->layerId();
  21952. if (!isBatchedAtLevel(grad_output, cur_level)) {
  21953. return at::_ops::_upsample_nearest_exact3d_backward::call(grad_output, output_size, input_size, scales_d, scales_h, scales_w);
  21954. }
  21955. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  21956. auto results = batch_rule(grad_output_value, grad_output_bdim, output_size, input_size, scales_d, scales_h, scales_w);
  21957. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21958. }
  21959. template <typename batch_rule_t, batch_rule_t batch_rule>
  21960. at::Tensor sigmoid_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & output) {
  21961. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21962. auto maybe_layer = maybeCurrentDynamicLayer();
  21963. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21964. int64_t cur_level = maybe_layer->layerId();
  21965. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(output, cur_level)) {
  21966. return at::_ops::sigmoid_backward::call(grad_output, output);
  21967. }
  21968. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  21969. auto [output_value, output_bdim] = unwrapTensorAtLevel(output, cur_level);
  21970. auto results = batch_rule(grad_output_value, grad_output_bdim, output_value, output_bdim);
  21971. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21972. }
  21973. template <typename batch_rule_t, batch_rule_t batch_rule>
  21974. at::Tensor logit_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self, ::std::optional<double> eps) {
  21975. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21976. auto maybe_layer = maybeCurrentDynamicLayer();
  21977. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21978. int64_t cur_level = maybe_layer->layerId();
  21979. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level)) {
  21980. return at::_ops::logit_backward::call(grad_output, self, eps);
  21981. }
  21982. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  21983. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  21984. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim, eps);
  21985. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  21986. }
  21987. template <typename batch_rule_t, batch_rule_t batch_rule>
  21988. at::Tensor tanh_backward_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & output) {
  21989. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  21990. auto maybe_layer = maybeCurrentDynamicLayer();
  21991. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  21992. int64_t cur_level = maybe_layer->layerId();
  21993. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(output, cur_level)) {
  21994. return at::_ops::tanh_backward::call(grad_output, output);
  21995. }
  21996. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  21997. auto [output_value, output_bdim] = unwrapTensorAtLevel(output, cur_level);
  21998. auto results = batch_rule(grad_output_value, grad_output_bdim, output_value, output_bdim);
  21999. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22000. }
  22001. template <typename batch_rule_t, batch_rule_t batch_rule>
  22002. at::Tensor slow_conv_transpose2d_generated_plumbing(const at::Tensor & self, const at::Tensor & weight, c10::SymIntArrayRef kernel_size, const ::std::optional<at::Tensor> & bias, c10::SymIntArrayRef stride, c10::SymIntArrayRef padding, c10::SymIntArrayRef output_padding, c10::SymIntArrayRef dilation) {
  22003. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22004. auto maybe_layer = maybeCurrentDynamicLayer();
  22005. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22006. int64_t cur_level = maybe_layer->layerId();
  22007. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  22008. return at::_ops::slow_conv_transpose2d::call(self, weight, kernel_size, bias, stride, padding, output_padding, dilation);
  22009. }
  22010. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22011. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  22012. std::optional<Tensor> bias_value;
  22013. std::optional<int64_t> bias_bdim;
  22014. if (bias) {
  22015. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  22016. }
  22017. auto results = batch_rule(self_value, self_bdim, weight_value, weight_bdim, kernel_size, bias_value, bias_bdim, stride, padding, output_padding, dilation);
  22018. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22019. }
  22020. template <typename batch_rule_t, batch_rule_t batch_rule>
  22021. at::Tensor slow_conv_transpose3d_generated_plumbing(const at::Tensor & self, const at::Tensor & weight, c10::SymIntArrayRef kernel_size, const ::std::optional<at::Tensor> & bias, c10::SymIntArrayRef stride, c10::SymIntArrayRef padding, c10::SymIntArrayRef output_padding, c10::SymIntArrayRef dilation) {
  22022. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22023. auto maybe_layer = maybeCurrentDynamicLayer();
  22024. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22025. int64_t cur_level = maybe_layer->layerId();
  22026. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  22027. return at::_ops::slow_conv_transpose3d::call(self, weight, kernel_size, bias, stride, padding, output_padding, dilation);
  22028. }
  22029. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22030. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  22031. std::optional<Tensor> bias_value;
  22032. std::optional<int64_t> bias_bdim;
  22033. if (bias) {
  22034. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  22035. }
  22036. auto results = batch_rule(self_value, self_bdim, weight_value, weight_bdim, kernel_size, bias_value, bias_bdim, stride, padding, output_padding, dilation);
  22037. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22038. }
  22039. template <typename batch_rule_t, batch_rule_t batch_rule>
  22040. at::Tensor thnn_conv2d_generated_plumbing(const at::Tensor & self, const at::Tensor & weight, c10::SymIntArrayRef kernel_size, const ::std::optional<at::Tensor> & bias, c10::SymIntArrayRef stride, c10::SymIntArrayRef padding) {
  22041. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22042. auto maybe_layer = maybeCurrentDynamicLayer();
  22043. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22044. int64_t cur_level = maybe_layer->layerId();
  22045. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  22046. return at::_ops::thnn_conv2d::call(self, weight, kernel_size, bias, stride, padding);
  22047. }
  22048. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22049. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  22050. std::optional<Tensor> bias_value;
  22051. std::optional<int64_t> bias_bdim;
  22052. if (bias) {
  22053. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  22054. }
  22055. auto results = batch_rule(self_value, self_bdim, weight_value, weight_bdim, kernel_size, bias_value, bias_bdim, stride, padding);
  22056. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22057. }
  22058. template <typename batch_rule_t, batch_rule_t batch_rule>
  22059. at::Tensor _slow_conv2d_forward_generated_plumbing(const at::Tensor & self, const at::Tensor & weight, c10::SymIntArrayRef kernel_size, const ::std::optional<at::Tensor> & bias, c10::SymIntArrayRef stride, c10::SymIntArrayRef padding) {
  22060. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22061. auto maybe_layer = maybeCurrentDynamicLayer();
  22062. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22063. int64_t cur_level = maybe_layer->layerId();
  22064. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  22065. return at::_ops::_slow_conv2d_forward::call(self, weight, kernel_size, bias, stride, padding);
  22066. }
  22067. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22068. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  22069. std::optional<Tensor> bias_value;
  22070. std::optional<int64_t> bias_bdim;
  22071. if (bias) {
  22072. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  22073. }
  22074. auto results = batch_rule(self_value, self_bdim, weight_value, weight_bdim, kernel_size, bias_value, bias_bdim, stride, padding);
  22075. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22076. }
  22077. template <typename batch_rule_t, batch_rule_t batch_rule>
  22078. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> _slow_conv2d_backward_output_mask_generated_plumbing(const at::Tensor & grad_output, const at::Tensor & self, const at::Tensor & weight, c10::SymIntArrayRef kernel_size, c10::SymIntArrayRef stride, c10::SymIntArrayRef padding, ::std::array<bool,3> output_mask) {
  22079. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22080. auto maybe_layer = maybeCurrentDynamicLayer();
  22081. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22082. int64_t cur_level = maybe_layer->layerId();
  22083. if (!isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(weight, cur_level)) {
  22084. return at::_ops::_slow_conv2d_backward_output_mask::call(grad_output, self, weight, kernel_size, stride, padding, output_mask);
  22085. }
  22086. auto [grad_output_value, grad_output_bdim] = unwrapTensorAtLevel(grad_output, cur_level);
  22087. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22088. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  22089. auto results = batch_rule(grad_output_value, grad_output_bdim, self_value, self_bdim, weight_value, weight_bdim, kernel_size, stride, padding, output_mask);
  22090. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  22091. }
  22092. template <typename batch_rule_t, batch_rule_t batch_rule>
  22093. at::Tensor _conv_depthwise2d_generated_plumbing(const at::Tensor & self, const at::Tensor & weight, c10::SymIntArrayRef kernel_size, const ::std::optional<at::Tensor> & bias, c10::SymIntArrayRef stride, c10::SymIntArrayRef padding, c10::SymIntArrayRef dilation) {
  22094. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22095. auto maybe_layer = maybeCurrentDynamicLayer();
  22096. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22097. int64_t cur_level = maybe_layer->layerId();
  22098. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  22099. return at::_ops::_conv_depthwise2d::call(self, weight, kernel_size, bias, stride, padding, dilation);
  22100. }
  22101. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22102. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  22103. std::optional<Tensor> bias_value;
  22104. std::optional<int64_t> bias_bdim;
  22105. if (bias) {
  22106. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  22107. }
  22108. auto results = batch_rule(self_value, self_bdim, weight_value, weight_bdim, kernel_size, bias_value, bias_bdim, stride, padding, dilation);
  22109. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22110. }
  22111. template <typename batch_rule_t, batch_rule_t batch_rule>
  22112. at::Tensor conv_depthwise3d_generated_plumbing(const at::Tensor & self, const at::Tensor & weight, c10::SymIntArrayRef kernel_size, const ::std::optional<at::Tensor> & bias, c10::SymIntArrayRef stride, c10::SymIntArrayRef padding, c10::SymIntArrayRef dilation) {
  22113. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22114. auto maybe_layer = maybeCurrentDynamicLayer();
  22115. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22116. int64_t cur_level = maybe_layer->layerId();
  22117. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  22118. return at::_ops::conv_depthwise3d::call(self, weight, kernel_size, bias, stride, padding, dilation);
  22119. }
  22120. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22121. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  22122. std::optional<Tensor> bias_value;
  22123. std::optional<int64_t> bias_bdim;
  22124. if (bias) {
  22125. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  22126. }
  22127. auto results = batch_rule(self_value, self_bdim, weight_value, weight_bdim, kernel_size, bias_value, bias_bdim, stride, padding, dilation);
  22128. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22129. }
  22130. template <typename batch_rule_t, batch_rule_t batch_rule>
  22131. at::Tensor slow_conv3d_generated_plumbing(const at::Tensor & self, const at::Tensor & weight, c10::SymIntArrayRef kernel_size, const ::std::optional<at::Tensor> & bias, c10::SymIntArrayRef stride, c10::SymIntArrayRef padding) {
  22132. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22133. auto maybe_layer = maybeCurrentDynamicLayer();
  22134. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22135. int64_t cur_level = maybe_layer->layerId();
  22136. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  22137. return at::_ops::slow_conv3d::call(self, weight, kernel_size, bias, stride, padding);
  22138. }
  22139. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22140. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  22141. std::optional<Tensor> bias_value;
  22142. std::optional<int64_t> bias_bdim;
  22143. if (bias) {
  22144. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  22145. }
  22146. auto results = batch_rule(self_value, self_bdim, weight_value, weight_bdim, kernel_size, bias_value, bias_bdim, stride, padding);
  22147. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22148. }
  22149. template <typename batch_rule_t, batch_rule_t batch_rule>
  22150. at::Tensor slow_conv3d_forward_generated_plumbing(const at::Tensor & self, const at::Tensor & weight, c10::SymIntArrayRef kernel_size, const ::std::optional<at::Tensor> & bias, c10::SymIntArrayRef stride, c10::SymIntArrayRef padding) {
  22151. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22152. auto maybe_layer = maybeCurrentDynamicLayer();
  22153. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22154. int64_t cur_level = maybe_layer->layerId();
  22155. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  22156. return at::_ops::slow_conv3d_forward::call(self, weight, kernel_size, bias, stride, padding);
  22157. }
  22158. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22159. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  22160. std::optional<Tensor> bias_value;
  22161. std::optional<int64_t> bias_bdim;
  22162. if (bias) {
  22163. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  22164. }
  22165. auto results = batch_rule(self_value, self_bdim, weight_value, weight_bdim, kernel_size, bias_value, bias_bdim, stride, padding);
  22166. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22167. }
  22168. template <typename batch_rule_t, batch_rule_t batch_rule>
  22169. at::Tensor slow_conv_dilated2d_generated_plumbing(const at::Tensor & self, const at::Tensor & weight, c10::SymIntArrayRef kernel_size, const ::std::optional<at::Tensor> & bias, c10::SymIntArrayRef stride, c10::SymIntArrayRef padding, c10::SymIntArrayRef dilation) {
  22170. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22171. auto maybe_layer = maybeCurrentDynamicLayer();
  22172. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22173. int64_t cur_level = maybe_layer->layerId();
  22174. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  22175. return at::_ops::slow_conv_dilated2d::call(self, weight, kernel_size, bias, stride, padding, dilation);
  22176. }
  22177. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22178. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  22179. std::optional<Tensor> bias_value;
  22180. std::optional<int64_t> bias_bdim;
  22181. if (bias) {
  22182. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  22183. }
  22184. auto results = batch_rule(self_value, self_bdim, weight_value, weight_bdim, kernel_size, bias_value, bias_bdim, stride, padding, dilation);
  22185. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22186. }
  22187. template <typename batch_rule_t, batch_rule_t batch_rule>
  22188. at::Tensor slow_conv_dilated3d_generated_plumbing(const at::Tensor & self, const at::Tensor & weight, c10::SymIntArrayRef kernel_size, const ::std::optional<at::Tensor> & bias, c10::SymIntArrayRef stride, c10::SymIntArrayRef padding, c10::SymIntArrayRef dilation) {
  22189. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22190. auto maybe_layer = maybeCurrentDynamicLayer();
  22191. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22192. int64_t cur_level = maybe_layer->layerId();
  22193. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level)) {
  22194. return at::_ops::slow_conv_dilated3d::call(self, weight, kernel_size, bias, stride, padding, dilation);
  22195. }
  22196. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22197. auto [weight_value, weight_bdim] = unwrapTensorAtLevel(weight, cur_level);
  22198. std::optional<Tensor> bias_value;
  22199. std::optional<int64_t> bias_bdim;
  22200. if (bias) {
  22201. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  22202. }
  22203. auto results = batch_rule(self_value, self_bdim, weight_value, weight_bdim, kernel_size, bias_value, bias_bdim, stride, padding, dilation);
  22204. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22205. }
  22206. template <typename batch_rule_t, batch_rule_t batch_rule>
  22207. at::Tensor col2im_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef output_size, at::IntArrayRef kernel_size, at::IntArrayRef dilation, at::IntArrayRef padding, at::IntArrayRef stride) {
  22208. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22209. auto maybe_layer = maybeCurrentDynamicLayer();
  22210. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22211. int64_t cur_level = maybe_layer->layerId();
  22212. if (!isBatchedAtLevel(self, cur_level)) {
  22213. return at::_ops::col2im::call(self, output_size, kernel_size, dilation, padding, stride);
  22214. }
  22215. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22216. auto results = batch_rule(self_value, self_bdim, output_size, kernel_size, dilation, padding, stride);
  22217. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22218. }
  22219. template <typename batch_rule_t, batch_rule_t batch_rule>
  22220. at::Tensor column_stack_generated_plumbing(at::TensorList tensors) {
  22221. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22222. auto maybe_layer = maybeCurrentDynamicLayer();
  22223. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22224. int64_t cur_level = maybe_layer->layerId();
  22225. if (!isBatchedAtLevel(tensors, cur_level)) {
  22226. return at::_ops::column_stack::call(tensors);
  22227. }
  22228. auto results = batch_rule(tensors);
  22229. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22230. }
  22231. template <typename batch_rule_t, batch_rule_t batch_rule>
  22232. at::Tensor im2col_generated_plumbing(const at::Tensor & self, at::IntArrayRef kernel_size, at::IntArrayRef dilation, at::IntArrayRef padding, at::IntArrayRef stride) {
  22233. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22234. auto maybe_layer = maybeCurrentDynamicLayer();
  22235. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22236. int64_t cur_level = maybe_layer->layerId();
  22237. if (!isBatchedAtLevel(self, cur_level)) {
  22238. return at::_ops::im2col::call(self, kernel_size, dilation, padding, stride);
  22239. }
  22240. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22241. auto results = batch_rule(self_value, self_bdim, kernel_size, dilation, padding, stride);
  22242. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22243. }
  22244. template <typename batch_rule_t, batch_rule_t batch_rule>
  22245. at::Tensor isfinite_generated_plumbing(const at::Tensor & self) {
  22246. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22247. auto maybe_layer = maybeCurrentDynamicLayer();
  22248. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22249. int64_t cur_level = maybe_layer->layerId();
  22250. if (!isBatchedAtLevel(self, cur_level)) {
  22251. return at::_ops::isfinite::call(self);
  22252. }
  22253. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22254. auto results = batch_rule(self_value, self_bdim);
  22255. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22256. }
  22257. template <typename batch_rule_t, batch_rule_t batch_rule>
  22258. at::Tensor isinf_generated_plumbing(const at::Tensor & self) {
  22259. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22260. auto maybe_layer = maybeCurrentDynamicLayer();
  22261. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22262. int64_t cur_level = maybe_layer->layerId();
  22263. if (!isBatchedAtLevel(self, cur_level)) {
  22264. return at::_ops::isinf::call(self);
  22265. }
  22266. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22267. auto results = batch_rule(self_value, self_bdim);
  22268. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22269. }
  22270. template <typename batch_rule_t, batch_rule_t batch_rule>
  22271. void record_stream_generated_plumbing(at::Tensor & self, at::Stream s) {
  22272. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22273. auto maybe_layer = maybeCurrentDynamicLayer();
  22274. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  22275. int64_t cur_level = maybe_layer->layerId();
  22276. if (!isBatchedAtLevel(self, cur_level)) {
  22277. return at::_ops::record_stream::call(self, s);
  22278. }
  22279. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22280. batch_rule(self_value, self_bdim, s);
  22281. }
  22282. template <typename batch_rule_t, batch_rule_t batch_rule>
  22283. at::Tensor isposinf_generated_plumbing(const at::Tensor & self) {
  22284. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22285. auto maybe_layer = maybeCurrentDynamicLayer();
  22286. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22287. int64_t cur_level = maybe_layer->layerId();
  22288. if (!isBatchedAtLevel(self, cur_level)) {
  22289. return at::_ops::isposinf::call(self);
  22290. }
  22291. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22292. auto results = batch_rule(self_value, self_bdim);
  22293. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22294. }
  22295. template <typename batch_rule_t, batch_rule_t batch_rule>
  22296. at::Tensor isneginf_generated_plumbing(const at::Tensor & self) {
  22297. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22298. auto maybe_layer = maybeCurrentDynamicLayer();
  22299. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22300. int64_t cur_level = maybe_layer->layerId();
  22301. if (!isBatchedAtLevel(self, cur_level)) {
  22302. return at::_ops::isneginf::call(self);
  22303. }
  22304. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22305. auto results = batch_rule(self_value, self_bdim);
  22306. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22307. }
  22308. template <typename batch_rule_t, batch_rule_t batch_rule>
  22309. at::Tensor _add_batch_dim_generated_plumbing(const at::Tensor & self, int64_t batch_dim, int64_t level) {
  22310. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22311. auto maybe_layer = maybeCurrentDynamicLayer();
  22312. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22313. int64_t cur_level = maybe_layer->layerId();
  22314. if (!isBatchedAtLevel(self, cur_level)) {
  22315. return at::_ops::_add_batch_dim::call(self, batch_dim, level);
  22316. }
  22317. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22318. auto results = batch_rule(self_value, self_bdim, batch_dim, level);
  22319. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22320. }
  22321. template <typename batch_rule_t, batch_rule_t batch_rule>
  22322. at::Tensor _remove_batch_dim_generated_plumbing(const at::Tensor & self, int64_t level, c10::SymInt batch_size, int64_t out_dim) {
  22323. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22324. auto maybe_layer = maybeCurrentDynamicLayer();
  22325. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22326. int64_t cur_level = maybe_layer->layerId();
  22327. if (!isBatchedAtLevel(self, cur_level)) {
  22328. return at::_ops::_remove_batch_dim::call(self, level, batch_size, out_dim);
  22329. }
  22330. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22331. auto results = batch_rule(self_value, self_bdim, level, batch_size, out_dim);
  22332. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22333. }
  22334. template <typename batch_rule_t, batch_rule_t batch_rule>
  22335. at::Tensor special_entr_generated_plumbing(const at::Tensor & self) {
  22336. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22337. auto maybe_layer = maybeCurrentDynamicLayer();
  22338. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22339. int64_t cur_level = maybe_layer->layerId();
  22340. if (!isBatchedAtLevel(self, cur_level)) {
  22341. return at::_ops::special_entr::call(self);
  22342. }
  22343. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22344. auto results = batch_rule(self_value, self_bdim);
  22345. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22346. }
  22347. template <typename batch_rule_t, batch_rule_t batch_rule>
  22348. at::Tensor special_ndtri_generated_plumbing(const at::Tensor & self) {
  22349. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22350. auto maybe_layer = maybeCurrentDynamicLayer();
  22351. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22352. int64_t cur_level = maybe_layer->layerId();
  22353. if (!isBatchedAtLevel(self, cur_level)) {
  22354. return at::_ops::special_ndtri::call(self);
  22355. }
  22356. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22357. auto results = batch_rule(self_value, self_bdim);
  22358. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22359. }
  22360. template <typename batch_rule_t, batch_rule_t batch_rule>
  22361. at::Tensor special_log_ndtr_generated_plumbing(const at::Tensor & self) {
  22362. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22363. auto maybe_layer = maybeCurrentDynamicLayer();
  22364. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22365. int64_t cur_level = maybe_layer->layerId();
  22366. if (!isBatchedAtLevel(self, cur_level)) {
  22367. return at::_ops::special_log_ndtr::call(self);
  22368. }
  22369. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22370. auto results = batch_rule(self_value, self_bdim);
  22371. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22372. }
  22373. template <typename batch_rule_t, batch_rule_t batch_rule>
  22374. at::Tensor special_expm1_generated_plumbing(const at::Tensor & self) {
  22375. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22376. auto maybe_layer = maybeCurrentDynamicLayer();
  22377. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22378. int64_t cur_level = maybe_layer->layerId();
  22379. if (!isBatchedAtLevel(self, cur_level)) {
  22380. return at::_ops::special_expm1::call(self);
  22381. }
  22382. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22383. auto results = batch_rule(self_value, self_bdim);
  22384. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22385. }
  22386. template <typename batch_rule_t, batch_rule_t batch_rule>
  22387. at::Tensor special_exp2_generated_plumbing(const at::Tensor & self) {
  22388. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22389. auto maybe_layer = maybeCurrentDynamicLayer();
  22390. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22391. int64_t cur_level = maybe_layer->layerId();
  22392. if (!isBatchedAtLevel(self, cur_level)) {
  22393. return at::_ops::special_exp2::call(self);
  22394. }
  22395. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22396. auto results = batch_rule(self_value, self_bdim);
  22397. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22398. }
  22399. template <typename batch_rule_t, batch_rule_t batch_rule>
  22400. at::Tensor special_psi_generated_plumbing(const at::Tensor & self) {
  22401. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22402. auto maybe_layer = maybeCurrentDynamicLayer();
  22403. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22404. int64_t cur_level = maybe_layer->layerId();
  22405. if (!isBatchedAtLevel(self, cur_level)) {
  22406. return at::_ops::special_psi::call(self);
  22407. }
  22408. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22409. auto results = batch_rule(self_value, self_bdim);
  22410. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22411. }
  22412. template <typename batch_rule_t, batch_rule_t batch_rule>
  22413. at::Tensor special_digamma_generated_plumbing(const at::Tensor & self) {
  22414. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22415. auto maybe_layer = maybeCurrentDynamicLayer();
  22416. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22417. int64_t cur_level = maybe_layer->layerId();
  22418. if (!isBatchedAtLevel(self, cur_level)) {
  22419. return at::_ops::special_digamma::call(self);
  22420. }
  22421. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22422. auto results = batch_rule(self_value, self_bdim);
  22423. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22424. }
  22425. template <typename batch_rule_t, batch_rule_t batch_rule>
  22426. at::Tensor special_gammaln_generated_plumbing(const at::Tensor & self) {
  22427. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22428. auto maybe_layer = maybeCurrentDynamicLayer();
  22429. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22430. int64_t cur_level = maybe_layer->layerId();
  22431. if (!isBatchedAtLevel(self, cur_level)) {
  22432. return at::_ops::special_gammaln::call(self);
  22433. }
  22434. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22435. auto results = batch_rule(self_value, self_bdim);
  22436. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22437. }
  22438. template <typename batch_rule_t, batch_rule_t batch_rule>
  22439. at::Tensor special_erf_generated_plumbing(const at::Tensor & self) {
  22440. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22441. auto maybe_layer = maybeCurrentDynamicLayer();
  22442. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22443. int64_t cur_level = maybe_layer->layerId();
  22444. if (!isBatchedAtLevel(self, cur_level)) {
  22445. return at::_ops::special_erf::call(self);
  22446. }
  22447. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22448. auto results = batch_rule(self_value, self_bdim);
  22449. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22450. }
  22451. template <typename batch_rule_t, batch_rule_t batch_rule>
  22452. at::Tensor special_erfc_generated_plumbing(const at::Tensor & self) {
  22453. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22454. auto maybe_layer = maybeCurrentDynamicLayer();
  22455. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22456. int64_t cur_level = maybe_layer->layerId();
  22457. if (!isBatchedAtLevel(self, cur_level)) {
  22458. return at::_ops::special_erfc::call(self);
  22459. }
  22460. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22461. auto results = batch_rule(self_value, self_bdim);
  22462. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22463. }
  22464. template <typename batch_rule_t, batch_rule_t batch_rule>
  22465. at::Tensor special_erfcx_generated_plumbing(const at::Tensor & self) {
  22466. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22467. auto maybe_layer = maybeCurrentDynamicLayer();
  22468. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22469. int64_t cur_level = maybe_layer->layerId();
  22470. if (!isBatchedAtLevel(self, cur_level)) {
  22471. return at::_ops::special_erfcx::call(self);
  22472. }
  22473. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22474. auto results = batch_rule(self_value, self_bdim);
  22475. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22476. }
  22477. template <typename batch_rule_t, batch_rule_t batch_rule>
  22478. at::Tensor special_erfinv_generated_plumbing(const at::Tensor & self) {
  22479. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22480. auto maybe_layer = maybeCurrentDynamicLayer();
  22481. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22482. int64_t cur_level = maybe_layer->layerId();
  22483. if (!isBatchedAtLevel(self, cur_level)) {
  22484. return at::_ops::special_erfinv::call(self);
  22485. }
  22486. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22487. auto results = batch_rule(self_value, self_bdim);
  22488. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22489. }
  22490. template <typename batch_rule_t, batch_rule_t batch_rule>
  22491. at::Tensor special_ndtr_generated_plumbing(const at::Tensor & self) {
  22492. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22493. auto maybe_layer = maybeCurrentDynamicLayer();
  22494. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22495. int64_t cur_level = maybe_layer->layerId();
  22496. if (!isBatchedAtLevel(self, cur_level)) {
  22497. return at::_ops::special_ndtr::call(self);
  22498. }
  22499. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22500. auto results = batch_rule(self_value, self_bdim);
  22501. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22502. }
  22503. template <typename batch_rule_t, batch_rule_t batch_rule>
  22504. at::Tensor special_xlog1py_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  22505. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22506. auto maybe_layer = maybeCurrentDynamicLayer();
  22507. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22508. int64_t cur_level = maybe_layer->layerId();
  22509. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  22510. return at::_ops::special_xlog1py::call(self, other);
  22511. }
  22512. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22513. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  22514. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  22515. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22516. }
  22517. template <typename batch_rule_t, batch_rule_t batch_rule>
  22518. at::Tensor special_xlog1py_self_scalar_generated_plumbing(const at::Scalar & self, const at::Tensor & other) {
  22519. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22520. auto maybe_layer = maybeCurrentDynamicLayer();
  22521. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22522. int64_t cur_level = maybe_layer->layerId();
  22523. if (!isBatchedAtLevel(other, cur_level)) {
  22524. return at::_ops::special_xlog1py_self_scalar::call(self, other);
  22525. }
  22526. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  22527. auto results = batch_rule(self, other_value, other_bdim);
  22528. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22529. }
  22530. template <typename batch_rule_t, batch_rule_t batch_rule>
  22531. at::Tensor special_xlog1py_other_scalar_generated_plumbing(const at::Tensor & self, const at::Scalar & other) {
  22532. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22533. auto maybe_layer = maybeCurrentDynamicLayer();
  22534. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22535. int64_t cur_level = maybe_layer->layerId();
  22536. if (!isBatchedAtLevel(self, cur_level)) {
  22537. return at::_ops::special_xlog1py_other_scalar::call(self, other);
  22538. }
  22539. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22540. auto results = batch_rule(self_value, self_bdim, other);
  22541. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22542. }
  22543. template <typename batch_rule_t, batch_rule_t batch_rule>
  22544. at::Tensor special_xlogy_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  22545. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22546. auto maybe_layer = maybeCurrentDynamicLayer();
  22547. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22548. int64_t cur_level = maybe_layer->layerId();
  22549. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  22550. return at::_ops::special_xlogy::call(self, other);
  22551. }
  22552. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22553. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  22554. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  22555. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22556. }
  22557. template <typename batch_rule_t, batch_rule_t batch_rule>
  22558. at::Tensor special_xlogy_self_scalar_generated_plumbing(const at::Scalar & self, const at::Tensor & other) {
  22559. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22560. auto maybe_layer = maybeCurrentDynamicLayer();
  22561. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22562. int64_t cur_level = maybe_layer->layerId();
  22563. if (!isBatchedAtLevel(other, cur_level)) {
  22564. return at::_ops::special_xlogy_self_scalar::call(self, other);
  22565. }
  22566. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  22567. auto results = batch_rule(self, other_value, other_bdim);
  22568. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22569. }
  22570. template <typename batch_rule_t, batch_rule_t batch_rule>
  22571. at::Tensor special_xlogy_other_scalar_generated_plumbing(const at::Tensor & self, const at::Scalar & other) {
  22572. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22573. auto maybe_layer = maybeCurrentDynamicLayer();
  22574. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22575. int64_t cur_level = maybe_layer->layerId();
  22576. if (!isBatchedAtLevel(self, cur_level)) {
  22577. return at::_ops::special_xlogy_other_scalar::call(self, other);
  22578. }
  22579. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22580. auto results = batch_rule(self_value, self_bdim, other);
  22581. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22582. }
  22583. template <typename batch_rule_t, batch_rule_t batch_rule>
  22584. at::Tensor special_zeta_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  22585. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22586. auto maybe_layer = maybeCurrentDynamicLayer();
  22587. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22588. int64_t cur_level = maybe_layer->layerId();
  22589. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  22590. return at::_ops::special_zeta::call(self, other);
  22591. }
  22592. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22593. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  22594. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  22595. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22596. }
  22597. template <typename batch_rule_t, batch_rule_t batch_rule>
  22598. at::Tensor special_zeta_self_scalar_generated_plumbing(const at::Scalar & self, const at::Tensor & other) {
  22599. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22600. auto maybe_layer = maybeCurrentDynamicLayer();
  22601. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22602. int64_t cur_level = maybe_layer->layerId();
  22603. if (!isBatchedAtLevel(other, cur_level)) {
  22604. return at::_ops::special_zeta_self_scalar::call(self, other);
  22605. }
  22606. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  22607. auto results = batch_rule(self, other_value, other_bdim);
  22608. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22609. }
  22610. template <typename batch_rule_t, batch_rule_t batch_rule>
  22611. at::Tensor special_zeta_other_scalar_generated_plumbing(const at::Tensor & self, const at::Scalar & other) {
  22612. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22613. auto maybe_layer = maybeCurrentDynamicLayer();
  22614. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22615. int64_t cur_level = maybe_layer->layerId();
  22616. if (!isBatchedAtLevel(self, cur_level)) {
  22617. return at::_ops::special_zeta_other_scalar::call(self, other);
  22618. }
  22619. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22620. auto results = batch_rule(self_value, self_bdim, other);
  22621. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22622. }
  22623. template <typename batch_rule_t, batch_rule_t batch_rule>
  22624. at::Tensor special_i0_generated_plumbing(const at::Tensor & self) {
  22625. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22626. auto maybe_layer = maybeCurrentDynamicLayer();
  22627. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22628. int64_t cur_level = maybe_layer->layerId();
  22629. if (!isBatchedAtLevel(self, cur_level)) {
  22630. return at::_ops::special_i0::call(self);
  22631. }
  22632. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22633. auto results = batch_rule(self_value, self_bdim);
  22634. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22635. }
  22636. template <typename batch_rule_t, batch_rule_t batch_rule>
  22637. at::Tensor special_i0e_generated_plumbing(const at::Tensor & self) {
  22638. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22639. auto maybe_layer = maybeCurrentDynamicLayer();
  22640. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22641. int64_t cur_level = maybe_layer->layerId();
  22642. if (!isBatchedAtLevel(self, cur_level)) {
  22643. return at::_ops::special_i0e::call(self);
  22644. }
  22645. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22646. auto results = batch_rule(self_value, self_bdim);
  22647. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22648. }
  22649. template <typename batch_rule_t, batch_rule_t batch_rule>
  22650. at::Tensor special_i1_generated_plumbing(const at::Tensor & self) {
  22651. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22652. auto maybe_layer = maybeCurrentDynamicLayer();
  22653. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22654. int64_t cur_level = maybe_layer->layerId();
  22655. if (!isBatchedAtLevel(self, cur_level)) {
  22656. return at::_ops::special_i1::call(self);
  22657. }
  22658. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22659. auto results = batch_rule(self_value, self_bdim);
  22660. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22661. }
  22662. template <typename batch_rule_t, batch_rule_t batch_rule>
  22663. at::Tensor special_i1e_generated_plumbing(const at::Tensor & self) {
  22664. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22665. auto maybe_layer = maybeCurrentDynamicLayer();
  22666. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22667. int64_t cur_level = maybe_layer->layerId();
  22668. if (!isBatchedAtLevel(self, cur_level)) {
  22669. return at::_ops::special_i1e::call(self);
  22670. }
  22671. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22672. auto results = batch_rule(self_value, self_bdim);
  22673. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22674. }
  22675. template <typename batch_rule_t, batch_rule_t batch_rule>
  22676. at::Tensor special_logit_generated_plumbing(const at::Tensor & self, ::std::optional<double> eps) {
  22677. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22678. auto maybe_layer = maybeCurrentDynamicLayer();
  22679. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22680. int64_t cur_level = maybe_layer->layerId();
  22681. if (!isBatchedAtLevel(self, cur_level)) {
  22682. return at::_ops::special_logit::call(self, eps);
  22683. }
  22684. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22685. auto results = batch_rule(self_value, self_bdim, eps);
  22686. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22687. }
  22688. template <typename batch_rule_t, batch_rule_t batch_rule>
  22689. at::Tensor special_polygamma_generated_plumbing(int64_t n, const at::Tensor & self) {
  22690. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22691. auto maybe_layer = maybeCurrentDynamicLayer();
  22692. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22693. int64_t cur_level = maybe_layer->layerId();
  22694. if (!isBatchedAtLevel(self, cur_level)) {
  22695. return at::_ops::special_polygamma::call(n, self);
  22696. }
  22697. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22698. auto results = batch_rule(n, self_value, self_bdim);
  22699. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22700. }
  22701. template <typename batch_rule_t, batch_rule_t batch_rule>
  22702. at::Tensor special_logsumexp_generated_plumbing(const at::Tensor & self, at::IntArrayRef dim, bool keepdim) {
  22703. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22704. auto maybe_layer = maybeCurrentDynamicLayer();
  22705. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22706. int64_t cur_level = maybe_layer->layerId();
  22707. if (!isBatchedAtLevel(self, cur_level)) {
  22708. return at::_ops::special_logsumexp::call(self, dim, keepdim);
  22709. }
  22710. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22711. auto results = batch_rule(self_value, self_bdim, dim, keepdim);
  22712. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22713. }
  22714. template <typename batch_rule_t, batch_rule_t batch_rule>
  22715. at::Tensor special_expit_generated_plumbing(const at::Tensor & self) {
  22716. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22717. auto maybe_layer = maybeCurrentDynamicLayer();
  22718. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22719. int64_t cur_level = maybe_layer->layerId();
  22720. if (!isBatchedAtLevel(self, cur_level)) {
  22721. return at::_ops::special_expit::call(self);
  22722. }
  22723. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22724. auto results = batch_rule(self_value, self_bdim);
  22725. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22726. }
  22727. template <typename batch_rule_t, batch_rule_t batch_rule>
  22728. at::Tensor special_sinc_generated_plumbing(const at::Tensor & self) {
  22729. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22730. auto maybe_layer = maybeCurrentDynamicLayer();
  22731. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22732. int64_t cur_level = maybe_layer->layerId();
  22733. if (!isBatchedAtLevel(self, cur_level)) {
  22734. return at::_ops::special_sinc::call(self);
  22735. }
  22736. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22737. auto results = batch_rule(self_value, self_bdim);
  22738. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22739. }
  22740. template <typename batch_rule_t, batch_rule_t batch_rule>
  22741. at::Tensor special_round_generated_plumbing(const at::Tensor & self, int64_t decimals) {
  22742. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22743. auto maybe_layer = maybeCurrentDynamicLayer();
  22744. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22745. int64_t cur_level = maybe_layer->layerId();
  22746. if (!isBatchedAtLevel(self, cur_level)) {
  22747. return at::_ops::special_round::call(self, decimals);
  22748. }
  22749. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22750. auto results = batch_rule(self_value, self_bdim, decimals);
  22751. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22752. }
  22753. template <typename batch_rule_t, batch_rule_t batch_rule>
  22754. at::Tensor special_log1p_generated_plumbing(const at::Tensor & self) {
  22755. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22756. auto maybe_layer = maybeCurrentDynamicLayer();
  22757. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22758. int64_t cur_level = maybe_layer->layerId();
  22759. if (!isBatchedAtLevel(self, cur_level)) {
  22760. return at::_ops::special_log1p::call(self);
  22761. }
  22762. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22763. auto results = batch_rule(self_value, self_bdim);
  22764. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22765. }
  22766. template <typename batch_rule_t, batch_rule_t batch_rule>
  22767. at::Tensor special_log_softmax_generated_plumbing(const at::Tensor & self, int64_t dim, ::std::optional<at::ScalarType> dtype) {
  22768. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22769. auto maybe_layer = maybeCurrentDynamicLayer();
  22770. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22771. int64_t cur_level = maybe_layer->layerId();
  22772. if (!isBatchedAtLevel(self, cur_level)) {
  22773. return at::_ops::special_log_softmax::call(self, dim, dtype);
  22774. }
  22775. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22776. auto results = batch_rule(self_value, self_bdim, dim, dtype);
  22777. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22778. }
  22779. template <typename batch_rule_t, batch_rule_t batch_rule>
  22780. at::Tensor special_gammainc_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  22781. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22782. auto maybe_layer = maybeCurrentDynamicLayer();
  22783. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22784. int64_t cur_level = maybe_layer->layerId();
  22785. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  22786. return at::_ops::special_gammainc::call(self, other);
  22787. }
  22788. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22789. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  22790. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  22791. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22792. }
  22793. template <typename batch_rule_t, batch_rule_t batch_rule>
  22794. at::Tensor special_gammaincc_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  22795. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22796. auto maybe_layer = maybeCurrentDynamicLayer();
  22797. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22798. int64_t cur_level = maybe_layer->layerId();
  22799. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  22800. return at::_ops::special_gammaincc::call(self, other);
  22801. }
  22802. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22803. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  22804. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  22805. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22806. }
  22807. template <typename batch_rule_t, batch_rule_t batch_rule>
  22808. at::Tensor special_multigammaln_generated_plumbing(const at::Tensor & self, int64_t p) {
  22809. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22810. auto maybe_layer = maybeCurrentDynamicLayer();
  22811. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22812. int64_t cur_level = maybe_layer->layerId();
  22813. if (!isBatchedAtLevel(self, cur_level)) {
  22814. return at::_ops::special_multigammaln::call(self, p);
  22815. }
  22816. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22817. auto results = batch_rule(self_value, self_bdim, p);
  22818. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22819. }
  22820. template <typename batch_rule_t, batch_rule_t batch_rule>
  22821. at::Tensor special_softmax_generated_plumbing(const at::Tensor & self, int64_t dim, ::std::optional<at::ScalarType> dtype) {
  22822. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22823. auto maybe_layer = maybeCurrentDynamicLayer();
  22824. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22825. int64_t cur_level = maybe_layer->layerId();
  22826. if (!isBatchedAtLevel(self, cur_level)) {
  22827. return at::_ops::special_softmax::call(self, dim, dtype);
  22828. }
  22829. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22830. auto results = batch_rule(self_value, self_bdim, dim, dtype);
  22831. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22832. }
  22833. template <typename batch_rule_t, batch_rule_t batch_rule>
  22834. at::Tensor fft_fft_generated_plumbing(const at::Tensor & self, ::std::optional<c10::SymInt> n, int64_t dim, ::std::optional<c10::string_view> norm) {
  22835. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22836. auto maybe_layer = maybeCurrentDynamicLayer();
  22837. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22838. int64_t cur_level = maybe_layer->layerId();
  22839. if (!isBatchedAtLevel(self, cur_level)) {
  22840. return at::_ops::fft_fft::call(self, n, dim, norm);
  22841. }
  22842. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22843. auto results = batch_rule(self_value, self_bdim, n, dim, norm);
  22844. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22845. }
  22846. template <typename batch_rule_t, batch_rule_t batch_rule>
  22847. at::Tensor fft_ifft_generated_plumbing(const at::Tensor & self, ::std::optional<c10::SymInt> n, int64_t dim, ::std::optional<c10::string_view> norm) {
  22848. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22849. auto maybe_layer = maybeCurrentDynamicLayer();
  22850. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22851. int64_t cur_level = maybe_layer->layerId();
  22852. if (!isBatchedAtLevel(self, cur_level)) {
  22853. return at::_ops::fft_ifft::call(self, n, dim, norm);
  22854. }
  22855. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22856. auto results = batch_rule(self_value, self_bdim, n, dim, norm);
  22857. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22858. }
  22859. template <typename batch_rule_t, batch_rule_t batch_rule>
  22860. at::Tensor fft_rfft_generated_plumbing(const at::Tensor & self, ::std::optional<c10::SymInt> n, int64_t dim, ::std::optional<c10::string_view> norm) {
  22861. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22862. auto maybe_layer = maybeCurrentDynamicLayer();
  22863. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22864. int64_t cur_level = maybe_layer->layerId();
  22865. if (!isBatchedAtLevel(self, cur_level)) {
  22866. return at::_ops::fft_rfft::call(self, n, dim, norm);
  22867. }
  22868. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22869. auto results = batch_rule(self_value, self_bdim, n, dim, norm);
  22870. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22871. }
  22872. template <typename batch_rule_t, batch_rule_t batch_rule>
  22873. at::Tensor fft_irfft_generated_plumbing(const at::Tensor & self, ::std::optional<c10::SymInt> n, int64_t dim, ::std::optional<c10::string_view> norm) {
  22874. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22875. auto maybe_layer = maybeCurrentDynamicLayer();
  22876. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22877. int64_t cur_level = maybe_layer->layerId();
  22878. if (!isBatchedAtLevel(self, cur_level)) {
  22879. return at::_ops::fft_irfft::call(self, n, dim, norm);
  22880. }
  22881. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22882. auto results = batch_rule(self_value, self_bdim, n, dim, norm);
  22883. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22884. }
  22885. template <typename batch_rule_t, batch_rule_t batch_rule>
  22886. at::Tensor fft_hfft_generated_plumbing(const at::Tensor & self, ::std::optional<c10::SymInt> n, int64_t dim, ::std::optional<c10::string_view> norm) {
  22887. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22888. auto maybe_layer = maybeCurrentDynamicLayer();
  22889. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22890. int64_t cur_level = maybe_layer->layerId();
  22891. if (!isBatchedAtLevel(self, cur_level)) {
  22892. return at::_ops::fft_hfft::call(self, n, dim, norm);
  22893. }
  22894. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22895. auto results = batch_rule(self_value, self_bdim, n, dim, norm);
  22896. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22897. }
  22898. template <typename batch_rule_t, batch_rule_t batch_rule>
  22899. at::Tensor fft_ihfft_generated_plumbing(const at::Tensor & self, ::std::optional<c10::SymInt> n, int64_t dim, ::std::optional<c10::string_view> norm) {
  22900. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22901. auto maybe_layer = maybeCurrentDynamicLayer();
  22902. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22903. int64_t cur_level = maybe_layer->layerId();
  22904. if (!isBatchedAtLevel(self, cur_level)) {
  22905. return at::_ops::fft_ihfft::call(self, n, dim, norm);
  22906. }
  22907. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22908. auto results = batch_rule(self_value, self_bdim, n, dim, norm);
  22909. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22910. }
  22911. template <typename batch_rule_t, batch_rule_t batch_rule>
  22912. at::Tensor fft_fft2_generated_plumbing(const at::Tensor & self, at::OptionalSymIntArrayRef s, at::IntArrayRef dim, ::std::optional<c10::string_view> norm) {
  22913. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22914. auto maybe_layer = maybeCurrentDynamicLayer();
  22915. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22916. int64_t cur_level = maybe_layer->layerId();
  22917. if (!isBatchedAtLevel(self, cur_level)) {
  22918. return at::_ops::fft_fft2::call(self, s, dim, norm);
  22919. }
  22920. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22921. auto results = batch_rule(self_value, self_bdim, s, dim, norm);
  22922. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22923. }
  22924. template <typename batch_rule_t, batch_rule_t batch_rule>
  22925. at::Tensor fft_ifft2_generated_plumbing(const at::Tensor & self, at::OptionalSymIntArrayRef s, at::IntArrayRef dim, ::std::optional<c10::string_view> norm) {
  22926. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22927. auto maybe_layer = maybeCurrentDynamicLayer();
  22928. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22929. int64_t cur_level = maybe_layer->layerId();
  22930. if (!isBatchedAtLevel(self, cur_level)) {
  22931. return at::_ops::fft_ifft2::call(self, s, dim, norm);
  22932. }
  22933. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22934. auto results = batch_rule(self_value, self_bdim, s, dim, norm);
  22935. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22936. }
  22937. template <typename batch_rule_t, batch_rule_t batch_rule>
  22938. at::Tensor fft_rfft2_generated_plumbing(const at::Tensor & self, at::OptionalSymIntArrayRef s, at::IntArrayRef dim, ::std::optional<c10::string_view> norm) {
  22939. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22940. auto maybe_layer = maybeCurrentDynamicLayer();
  22941. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22942. int64_t cur_level = maybe_layer->layerId();
  22943. if (!isBatchedAtLevel(self, cur_level)) {
  22944. return at::_ops::fft_rfft2::call(self, s, dim, norm);
  22945. }
  22946. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22947. auto results = batch_rule(self_value, self_bdim, s, dim, norm);
  22948. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22949. }
  22950. template <typename batch_rule_t, batch_rule_t batch_rule>
  22951. at::Tensor fft_irfft2_generated_plumbing(const at::Tensor & self, at::OptionalSymIntArrayRef s, at::IntArrayRef dim, ::std::optional<c10::string_view> norm) {
  22952. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22953. auto maybe_layer = maybeCurrentDynamicLayer();
  22954. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22955. int64_t cur_level = maybe_layer->layerId();
  22956. if (!isBatchedAtLevel(self, cur_level)) {
  22957. return at::_ops::fft_irfft2::call(self, s, dim, norm);
  22958. }
  22959. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22960. auto results = batch_rule(self_value, self_bdim, s, dim, norm);
  22961. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22962. }
  22963. template <typename batch_rule_t, batch_rule_t batch_rule>
  22964. at::Tensor fft_hfft2_generated_plumbing(const at::Tensor & self, at::OptionalSymIntArrayRef s, at::IntArrayRef dim, ::std::optional<c10::string_view> norm) {
  22965. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22966. auto maybe_layer = maybeCurrentDynamicLayer();
  22967. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22968. int64_t cur_level = maybe_layer->layerId();
  22969. if (!isBatchedAtLevel(self, cur_level)) {
  22970. return at::_ops::fft_hfft2::call(self, s, dim, norm);
  22971. }
  22972. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22973. auto results = batch_rule(self_value, self_bdim, s, dim, norm);
  22974. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22975. }
  22976. template <typename batch_rule_t, batch_rule_t batch_rule>
  22977. at::Tensor fft_ihfft2_generated_plumbing(const at::Tensor & self, at::OptionalSymIntArrayRef s, at::IntArrayRef dim, ::std::optional<c10::string_view> norm) {
  22978. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22979. auto maybe_layer = maybeCurrentDynamicLayer();
  22980. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22981. int64_t cur_level = maybe_layer->layerId();
  22982. if (!isBatchedAtLevel(self, cur_level)) {
  22983. return at::_ops::fft_ihfft2::call(self, s, dim, norm);
  22984. }
  22985. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22986. auto results = batch_rule(self_value, self_bdim, s, dim, norm);
  22987. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  22988. }
  22989. template <typename batch_rule_t, batch_rule_t batch_rule>
  22990. at::Tensor fft_fftn_generated_plumbing(const at::Tensor & self, at::OptionalSymIntArrayRef s, at::OptionalIntArrayRef dim, ::std::optional<c10::string_view> norm) {
  22991. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  22992. auto maybe_layer = maybeCurrentDynamicLayer();
  22993. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  22994. int64_t cur_level = maybe_layer->layerId();
  22995. if (!isBatchedAtLevel(self, cur_level)) {
  22996. return at::_ops::fft_fftn::call(self, s, dim, norm);
  22997. }
  22998. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  22999. auto results = batch_rule(self_value, self_bdim, s, dim, norm);
  23000. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23001. }
  23002. template <typename batch_rule_t, batch_rule_t batch_rule>
  23003. at::Tensor fft_ifftn_generated_plumbing(const at::Tensor & self, at::OptionalSymIntArrayRef s, at::OptionalIntArrayRef dim, ::std::optional<c10::string_view> norm) {
  23004. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23005. auto maybe_layer = maybeCurrentDynamicLayer();
  23006. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23007. int64_t cur_level = maybe_layer->layerId();
  23008. if (!isBatchedAtLevel(self, cur_level)) {
  23009. return at::_ops::fft_ifftn::call(self, s, dim, norm);
  23010. }
  23011. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23012. auto results = batch_rule(self_value, self_bdim, s, dim, norm);
  23013. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23014. }
  23015. template <typename batch_rule_t, batch_rule_t batch_rule>
  23016. at::Tensor fft_rfftn_generated_plumbing(const at::Tensor & self, at::OptionalSymIntArrayRef s, at::OptionalIntArrayRef dim, ::std::optional<c10::string_view> norm) {
  23017. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23018. auto maybe_layer = maybeCurrentDynamicLayer();
  23019. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23020. int64_t cur_level = maybe_layer->layerId();
  23021. if (!isBatchedAtLevel(self, cur_level)) {
  23022. return at::_ops::fft_rfftn::call(self, s, dim, norm);
  23023. }
  23024. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23025. auto results = batch_rule(self_value, self_bdim, s, dim, norm);
  23026. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23027. }
  23028. template <typename batch_rule_t, batch_rule_t batch_rule>
  23029. at::Tensor fft_irfftn_generated_plumbing(const at::Tensor & self, at::OptionalSymIntArrayRef s, at::OptionalIntArrayRef dim, ::std::optional<c10::string_view> norm) {
  23030. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23031. auto maybe_layer = maybeCurrentDynamicLayer();
  23032. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23033. int64_t cur_level = maybe_layer->layerId();
  23034. if (!isBatchedAtLevel(self, cur_level)) {
  23035. return at::_ops::fft_irfftn::call(self, s, dim, norm);
  23036. }
  23037. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23038. auto results = batch_rule(self_value, self_bdim, s, dim, norm);
  23039. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23040. }
  23041. template <typename batch_rule_t, batch_rule_t batch_rule>
  23042. at::Tensor fft_hfftn_generated_plumbing(const at::Tensor & self, at::OptionalSymIntArrayRef s, at::OptionalIntArrayRef dim, ::std::optional<c10::string_view> norm) {
  23043. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23044. auto maybe_layer = maybeCurrentDynamicLayer();
  23045. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23046. int64_t cur_level = maybe_layer->layerId();
  23047. if (!isBatchedAtLevel(self, cur_level)) {
  23048. return at::_ops::fft_hfftn::call(self, s, dim, norm);
  23049. }
  23050. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23051. auto results = batch_rule(self_value, self_bdim, s, dim, norm);
  23052. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23053. }
  23054. template <typename batch_rule_t, batch_rule_t batch_rule>
  23055. at::Tensor fft_ihfftn_generated_plumbing(const at::Tensor & self, at::OptionalSymIntArrayRef s, at::OptionalIntArrayRef dim, ::std::optional<c10::string_view> norm) {
  23056. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23057. auto maybe_layer = maybeCurrentDynamicLayer();
  23058. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23059. int64_t cur_level = maybe_layer->layerId();
  23060. if (!isBatchedAtLevel(self, cur_level)) {
  23061. return at::_ops::fft_ihfftn::call(self, s, dim, norm);
  23062. }
  23063. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23064. auto results = batch_rule(self_value, self_bdim, s, dim, norm);
  23065. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23066. }
  23067. template <typename batch_rule_t, batch_rule_t batch_rule>
  23068. at::Tensor fft_fftshift_generated_plumbing(const at::Tensor & self, at::OptionalIntArrayRef dim) {
  23069. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23070. auto maybe_layer = maybeCurrentDynamicLayer();
  23071. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23072. int64_t cur_level = maybe_layer->layerId();
  23073. if (!isBatchedAtLevel(self, cur_level)) {
  23074. return at::_ops::fft_fftshift::call(self, dim);
  23075. }
  23076. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23077. auto results = batch_rule(self_value, self_bdim, dim);
  23078. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23079. }
  23080. template <typename batch_rule_t, batch_rule_t batch_rule>
  23081. at::Tensor fft_ifftshift_generated_plumbing(const at::Tensor & self, at::OptionalIntArrayRef dim) {
  23082. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23083. auto maybe_layer = maybeCurrentDynamicLayer();
  23084. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23085. int64_t cur_level = maybe_layer->layerId();
  23086. if (!isBatchedAtLevel(self, cur_level)) {
  23087. return at::_ops::fft_ifftshift::call(self, dim);
  23088. }
  23089. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23090. auto results = batch_rule(self_value, self_bdim, dim);
  23091. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23092. }
  23093. template <typename batch_rule_t, batch_rule_t batch_rule>
  23094. ::std::tuple<at::Tensor,at::Tensor> linalg_cholesky_ex_generated_plumbing(const at::Tensor & self, bool upper, bool check_errors) {
  23095. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23096. auto maybe_layer = maybeCurrentDynamicLayer();
  23097. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23098. int64_t cur_level = maybe_layer->layerId();
  23099. if (!isBatchedAtLevel(self, cur_level)) {
  23100. return at::_ops::linalg_cholesky_ex::call(self, upper, check_errors);
  23101. }
  23102. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23103. auto results = batch_rule(self_value, self_bdim, upper, check_errors);
  23104. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  23105. }
  23106. template <typename batch_rule_t, batch_rule_t batch_rule>
  23107. at::Tensor linalg_cholesky_generated_plumbing(const at::Tensor & self, bool upper) {
  23108. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23109. auto maybe_layer = maybeCurrentDynamicLayer();
  23110. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23111. int64_t cur_level = maybe_layer->layerId();
  23112. if (!isBatchedAtLevel(self, cur_level)) {
  23113. return at::_ops::linalg_cholesky::call(self, upper);
  23114. }
  23115. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23116. auto results = batch_rule(self_value, self_bdim, upper);
  23117. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23118. }
  23119. template <typename batch_rule_t, batch_rule_t batch_rule>
  23120. at::Tensor linalg_cross_generated_plumbing(const at::Tensor & self, const at::Tensor & other, int64_t dim) {
  23121. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23122. auto maybe_layer = maybeCurrentDynamicLayer();
  23123. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23124. int64_t cur_level = maybe_layer->layerId();
  23125. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  23126. return at::_ops::linalg_cross::call(self, other, dim);
  23127. }
  23128. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23129. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  23130. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim, dim);
  23131. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23132. }
  23133. template <typename batch_rule_t, batch_rule_t batch_rule>
  23134. ::std::tuple<at::Tensor,at::Tensor> linalg_lu_factor_generated_plumbing(const at::Tensor & A, bool pivot) {
  23135. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23136. auto maybe_layer = maybeCurrentDynamicLayer();
  23137. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23138. int64_t cur_level = maybe_layer->layerId();
  23139. if (!isBatchedAtLevel(A, cur_level)) {
  23140. return at::_ops::linalg_lu_factor::call(A, pivot);
  23141. }
  23142. auto [A_value, A_bdim] = unwrapTensorAtLevel(A, cur_level);
  23143. auto results = batch_rule(A_value, A_bdim, pivot);
  23144. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  23145. }
  23146. template <typename batch_rule_t, batch_rule_t batch_rule>
  23147. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> linalg_lu_factor_ex_generated_plumbing(const at::Tensor & A, bool pivot, bool check_errors) {
  23148. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23149. auto maybe_layer = maybeCurrentDynamicLayer();
  23150. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23151. int64_t cur_level = maybe_layer->layerId();
  23152. if (!isBatchedAtLevel(A, cur_level)) {
  23153. return at::_ops::linalg_lu_factor_ex::call(A, pivot, check_errors);
  23154. }
  23155. auto [A_value, A_bdim] = unwrapTensorAtLevel(A, cur_level);
  23156. auto results = batch_rule(A_value, A_bdim, pivot, check_errors);
  23157. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  23158. }
  23159. template <typename batch_rule_t, batch_rule_t batch_rule>
  23160. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> linalg_lu_generated_plumbing(const at::Tensor & A, bool pivot) {
  23161. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23162. auto maybe_layer = maybeCurrentDynamicLayer();
  23163. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23164. int64_t cur_level = maybe_layer->layerId();
  23165. if (!isBatchedAtLevel(A, cur_level)) {
  23166. return at::_ops::linalg_lu::call(A, pivot);
  23167. }
  23168. auto [A_value, A_bdim] = unwrapTensorAtLevel(A, cur_level);
  23169. auto results = batch_rule(A_value, A_bdim, pivot);
  23170. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  23171. }
  23172. template <typename batch_rule_t, batch_rule_t batch_rule>
  23173. at::Tensor linalg_lu_solve_generated_plumbing(const at::Tensor & LU, const at::Tensor & pivots, const at::Tensor & B, bool left, bool adjoint) {
  23174. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23175. auto maybe_layer = maybeCurrentDynamicLayer();
  23176. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23177. int64_t cur_level = maybe_layer->layerId();
  23178. if (!isBatchedAtLevel(LU, cur_level) && !isBatchedAtLevel(pivots, cur_level) && !isBatchedAtLevel(B, cur_level)) {
  23179. return at::_ops::linalg_lu_solve::call(LU, pivots, B, left, adjoint);
  23180. }
  23181. auto [LU_value, LU_bdim] = unwrapTensorAtLevel(LU, cur_level);
  23182. auto [pivots_value, pivots_bdim] = unwrapTensorAtLevel(pivots, cur_level);
  23183. auto [B_value, B_bdim] = unwrapTensorAtLevel(B, cur_level);
  23184. auto results = batch_rule(LU_value, LU_bdim, pivots_value, pivots_bdim, B_value, B_bdim, left, adjoint);
  23185. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23186. }
  23187. template <typename batch_rule_t, batch_rule_t batch_rule>
  23188. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> _linalg_det_generated_plumbing(const at::Tensor & A) {
  23189. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23190. auto maybe_layer = maybeCurrentDynamicLayer();
  23191. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23192. int64_t cur_level = maybe_layer->layerId();
  23193. if (!isBatchedAtLevel(A, cur_level)) {
  23194. return at::_ops::_linalg_det::call(A);
  23195. }
  23196. auto [A_value, A_bdim] = unwrapTensorAtLevel(A, cur_level);
  23197. auto results = batch_rule(A_value, A_bdim);
  23198. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  23199. }
  23200. template <typename batch_rule_t, batch_rule_t batch_rule>
  23201. at::Tensor linalg_det_generated_plumbing(const at::Tensor & A) {
  23202. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23203. auto maybe_layer = maybeCurrentDynamicLayer();
  23204. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23205. int64_t cur_level = maybe_layer->layerId();
  23206. if (!isBatchedAtLevel(A, cur_level)) {
  23207. return at::_ops::linalg_det::call(A);
  23208. }
  23209. auto [A_value, A_bdim] = unwrapTensorAtLevel(A, cur_level);
  23210. auto results = batch_rule(A_value, A_bdim);
  23211. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23212. }
  23213. template <typename batch_rule_t, batch_rule_t batch_rule>
  23214. at::Tensor det_generated_plumbing(const at::Tensor & self) {
  23215. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23216. auto maybe_layer = maybeCurrentDynamicLayer();
  23217. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23218. int64_t cur_level = maybe_layer->layerId();
  23219. if (!isBatchedAtLevel(self, cur_level)) {
  23220. return at::_ops::det::call(self);
  23221. }
  23222. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23223. auto results = batch_rule(self_value, self_bdim);
  23224. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23225. }
  23226. template <typename batch_rule_t, batch_rule_t batch_rule>
  23227. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> linalg_ldl_factor_ex_generated_plumbing(const at::Tensor & self, bool hermitian, bool check_errors) {
  23228. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23229. auto maybe_layer = maybeCurrentDynamicLayer();
  23230. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23231. int64_t cur_level = maybe_layer->layerId();
  23232. if (!isBatchedAtLevel(self, cur_level)) {
  23233. return at::_ops::linalg_ldl_factor_ex::call(self, hermitian, check_errors);
  23234. }
  23235. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23236. auto results = batch_rule(self_value, self_bdim, hermitian, check_errors);
  23237. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  23238. }
  23239. template <typename batch_rule_t, batch_rule_t batch_rule>
  23240. ::std::tuple<at::Tensor,at::Tensor> linalg_ldl_factor_generated_plumbing(const at::Tensor & self, bool hermitian) {
  23241. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23242. auto maybe_layer = maybeCurrentDynamicLayer();
  23243. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23244. int64_t cur_level = maybe_layer->layerId();
  23245. if (!isBatchedAtLevel(self, cur_level)) {
  23246. return at::_ops::linalg_ldl_factor::call(self, hermitian);
  23247. }
  23248. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23249. auto results = batch_rule(self_value, self_bdim, hermitian);
  23250. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  23251. }
  23252. template <typename batch_rule_t, batch_rule_t batch_rule>
  23253. at::Tensor linalg_ldl_solve_generated_plumbing(const at::Tensor & LD, const at::Tensor & pivots, const at::Tensor & B, bool hermitian) {
  23254. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23255. auto maybe_layer = maybeCurrentDynamicLayer();
  23256. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23257. int64_t cur_level = maybe_layer->layerId();
  23258. if (!isBatchedAtLevel(LD, cur_level) && !isBatchedAtLevel(pivots, cur_level) && !isBatchedAtLevel(B, cur_level)) {
  23259. return at::_ops::linalg_ldl_solve::call(LD, pivots, B, hermitian);
  23260. }
  23261. auto [LD_value, LD_bdim] = unwrapTensorAtLevel(LD, cur_level);
  23262. auto [pivots_value, pivots_bdim] = unwrapTensorAtLevel(pivots, cur_level);
  23263. auto [B_value, B_bdim] = unwrapTensorAtLevel(B, cur_level);
  23264. auto results = batch_rule(LD_value, LD_bdim, pivots_value, pivots_bdim, B_value, B_bdim, hermitian);
  23265. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23266. }
  23267. template <typename batch_rule_t, batch_rule_t batch_rule>
  23268. ::std::tuple<at::Tensor,at::Tensor,at::Tensor,at::Tensor> linalg_lstsq_generated_plumbing(const at::Tensor & self, const at::Tensor & b, ::std::optional<double> rcond, ::std::optional<c10::string_view> driver) {
  23269. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23270. auto maybe_layer = maybeCurrentDynamicLayer();
  23271. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23272. int64_t cur_level = maybe_layer->layerId();
  23273. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(b, cur_level)) {
  23274. return at::_ops::linalg_lstsq::call(self, b, rcond, driver);
  23275. }
  23276. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23277. auto [b_value, b_bdim] = unwrapTensorAtLevel(b, cur_level);
  23278. auto results = batch_rule(self_value, self_bdim, b_value, b_bdim, rcond, driver);
  23279. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level), makeBatched(std::get<6>(results), std::get<7>(results), cur_level));
  23280. }
  23281. template <typename batch_rule_t, batch_rule_t batch_rule>
  23282. at::Tensor linalg_matmul_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  23283. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23284. auto maybe_layer = maybeCurrentDynamicLayer();
  23285. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23286. int64_t cur_level = maybe_layer->layerId();
  23287. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  23288. return at::_ops::linalg_matmul::call(self, other);
  23289. }
  23290. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23291. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  23292. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  23293. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23294. }
  23295. template <typename batch_rule_t, batch_rule_t batch_rule>
  23296. at::Tensor linalg_vecdot_generated_plumbing(const at::Tensor & x, const at::Tensor & y, int64_t dim) {
  23297. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23298. auto maybe_layer = maybeCurrentDynamicLayer();
  23299. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23300. int64_t cur_level = maybe_layer->layerId();
  23301. if (!isBatchedAtLevel(x, cur_level) && !isBatchedAtLevel(y, cur_level)) {
  23302. return at::_ops::linalg_vecdot::call(x, y, dim);
  23303. }
  23304. auto [x_value, x_bdim] = unwrapTensorAtLevel(x, cur_level);
  23305. auto [y_value, y_bdim] = unwrapTensorAtLevel(y, cur_level);
  23306. auto results = batch_rule(x_value, x_bdim, y_value, y_bdim, dim);
  23307. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23308. }
  23309. template <typename batch_rule_t, batch_rule_t batch_rule>
  23310. at::Tensor linalg_matrix_exp_generated_plumbing(const at::Tensor & self) {
  23311. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23312. auto maybe_layer = maybeCurrentDynamicLayer();
  23313. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23314. int64_t cur_level = maybe_layer->layerId();
  23315. if (!isBatchedAtLevel(self, cur_level)) {
  23316. return at::_ops::linalg_matrix_exp::call(self);
  23317. }
  23318. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23319. auto results = batch_rule(self_value, self_bdim);
  23320. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23321. }
  23322. template <typename batch_rule_t, batch_rule_t batch_rule>
  23323. ::std::tuple<at::Tensor,at::Tensor,at::Tensor,at::Tensor> _linalg_slogdet_generated_plumbing(const at::Tensor & A) {
  23324. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23325. auto maybe_layer = maybeCurrentDynamicLayer();
  23326. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23327. int64_t cur_level = maybe_layer->layerId();
  23328. if (!isBatchedAtLevel(A, cur_level)) {
  23329. return at::_ops::_linalg_slogdet::call(A);
  23330. }
  23331. auto [A_value, A_bdim] = unwrapTensorAtLevel(A, cur_level);
  23332. auto results = batch_rule(A_value, A_bdim);
  23333. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level), makeBatched(std::get<6>(results), std::get<7>(results), cur_level));
  23334. }
  23335. template <typename batch_rule_t, batch_rule_t batch_rule>
  23336. ::std::tuple<at::Tensor,at::Tensor> linalg_slogdet_generated_plumbing(const at::Tensor & A) {
  23337. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23338. auto maybe_layer = maybeCurrentDynamicLayer();
  23339. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23340. int64_t cur_level = maybe_layer->layerId();
  23341. if (!isBatchedAtLevel(A, cur_level)) {
  23342. return at::_ops::linalg_slogdet::call(A);
  23343. }
  23344. auto [A_value, A_bdim] = unwrapTensorAtLevel(A, cur_level);
  23345. auto results = batch_rule(A_value, A_bdim);
  23346. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  23347. }
  23348. template <typename batch_rule_t, batch_rule_t batch_rule>
  23349. ::std::tuple<at::Tensor,at::Tensor> slogdet_generated_plumbing(const at::Tensor & self) {
  23350. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23351. auto maybe_layer = maybeCurrentDynamicLayer();
  23352. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23353. int64_t cur_level = maybe_layer->layerId();
  23354. if (!isBatchedAtLevel(self, cur_level)) {
  23355. return at::_ops::slogdet::call(self);
  23356. }
  23357. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23358. auto results = batch_rule(self_value, self_bdim);
  23359. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  23360. }
  23361. template <typename batch_rule_t, batch_rule_t batch_rule>
  23362. at::Tensor logdet_generated_plumbing(const at::Tensor & self) {
  23363. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23364. auto maybe_layer = maybeCurrentDynamicLayer();
  23365. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23366. int64_t cur_level = maybe_layer->layerId();
  23367. if (!isBatchedAtLevel(self, cur_level)) {
  23368. return at::_ops::logdet::call(self);
  23369. }
  23370. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23371. auto results = batch_rule(self_value, self_bdim);
  23372. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23373. }
  23374. template <typename batch_rule_t, batch_rule_t batch_rule>
  23375. ::std::tuple<at::Tensor,at::Tensor> linalg_eig_generated_plumbing(const at::Tensor & self) {
  23376. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23377. auto maybe_layer = maybeCurrentDynamicLayer();
  23378. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23379. int64_t cur_level = maybe_layer->layerId();
  23380. if (!isBatchedAtLevel(self, cur_level)) {
  23381. return at::_ops::linalg_eig::call(self);
  23382. }
  23383. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23384. auto results = batch_rule(self_value, self_bdim);
  23385. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  23386. }
  23387. template <typename batch_rule_t, batch_rule_t batch_rule>
  23388. at::Tensor _linalg_eigvals_generated_plumbing(const at::Tensor & self) {
  23389. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23390. auto maybe_layer = maybeCurrentDynamicLayer();
  23391. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23392. int64_t cur_level = maybe_layer->layerId();
  23393. if (!isBatchedAtLevel(self, cur_level)) {
  23394. return at::_ops::_linalg_eigvals::call(self);
  23395. }
  23396. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23397. auto results = batch_rule(self_value, self_bdim);
  23398. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23399. }
  23400. template <typename batch_rule_t, batch_rule_t batch_rule>
  23401. at::Tensor linalg_eigvals_generated_plumbing(const at::Tensor & self) {
  23402. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23403. auto maybe_layer = maybeCurrentDynamicLayer();
  23404. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23405. int64_t cur_level = maybe_layer->layerId();
  23406. if (!isBatchedAtLevel(self, cur_level)) {
  23407. return at::_ops::linalg_eigvals::call(self);
  23408. }
  23409. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23410. auto results = batch_rule(self_value, self_bdim);
  23411. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23412. }
  23413. template <typename batch_rule_t, batch_rule_t batch_rule>
  23414. ::std::tuple<at::Tensor,at::Tensor> _linalg_eigh_generated_plumbing(const at::Tensor & A, c10::string_view UPLO, bool compute_v) {
  23415. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23416. auto maybe_layer = maybeCurrentDynamicLayer();
  23417. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23418. int64_t cur_level = maybe_layer->layerId();
  23419. if (!isBatchedAtLevel(A, cur_level)) {
  23420. return at::_ops::_linalg_eigh::call(A, UPLO, compute_v);
  23421. }
  23422. auto [A_value, A_bdim] = unwrapTensorAtLevel(A, cur_level);
  23423. auto results = batch_rule(A_value, A_bdim, UPLO, compute_v);
  23424. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  23425. }
  23426. template <typename batch_rule_t, batch_rule_t batch_rule>
  23427. ::std::tuple<at::Tensor,at::Tensor> linalg_eigh_generated_plumbing(const at::Tensor & self, c10::string_view UPLO) {
  23428. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23429. auto maybe_layer = maybeCurrentDynamicLayer();
  23430. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23431. int64_t cur_level = maybe_layer->layerId();
  23432. if (!isBatchedAtLevel(self, cur_level)) {
  23433. return at::_ops::linalg_eigh::call(self, UPLO);
  23434. }
  23435. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23436. auto results = batch_rule(self_value, self_bdim, UPLO);
  23437. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  23438. }
  23439. template <typename batch_rule_t, batch_rule_t batch_rule>
  23440. at::Tensor linalg_eigvalsh_generated_plumbing(const at::Tensor & self, c10::string_view UPLO) {
  23441. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23442. auto maybe_layer = maybeCurrentDynamicLayer();
  23443. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23444. int64_t cur_level = maybe_layer->layerId();
  23445. if (!isBatchedAtLevel(self, cur_level)) {
  23446. return at::_ops::linalg_eigvalsh::call(self, UPLO);
  23447. }
  23448. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23449. auto results = batch_rule(self_value, self_bdim, UPLO);
  23450. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23451. }
  23452. template <typename batch_rule_t, batch_rule_t batch_rule>
  23453. at::Tensor linalg_householder_product_generated_plumbing(const at::Tensor & input, const at::Tensor & tau) {
  23454. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23455. auto maybe_layer = maybeCurrentDynamicLayer();
  23456. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23457. int64_t cur_level = maybe_layer->layerId();
  23458. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(tau, cur_level)) {
  23459. return at::_ops::linalg_householder_product::call(input, tau);
  23460. }
  23461. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  23462. auto [tau_value, tau_bdim] = unwrapTensorAtLevel(tau, cur_level);
  23463. auto results = batch_rule(input_value, input_bdim, tau_value, tau_bdim);
  23464. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23465. }
  23466. template <typename batch_rule_t, batch_rule_t batch_rule>
  23467. ::std::tuple<at::Tensor,at::Tensor> linalg_inv_ex_generated_plumbing(const at::Tensor & A, bool check_errors) {
  23468. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23469. auto maybe_layer = maybeCurrentDynamicLayer();
  23470. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23471. int64_t cur_level = maybe_layer->layerId();
  23472. if (!isBatchedAtLevel(A, cur_level)) {
  23473. return at::_ops::linalg_inv_ex::call(A, check_errors);
  23474. }
  23475. auto [A_value, A_bdim] = unwrapTensorAtLevel(A, cur_level);
  23476. auto results = batch_rule(A_value, A_bdim, check_errors);
  23477. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  23478. }
  23479. template <typename batch_rule_t, batch_rule_t batch_rule>
  23480. at::Tensor linalg_inv_generated_plumbing(const at::Tensor & A) {
  23481. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23482. auto maybe_layer = maybeCurrentDynamicLayer();
  23483. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23484. int64_t cur_level = maybe_layer->layerId();
  23485. if (!isBatchedAtLevel(A, cur_level)) {
  23486. return at::_ops::linalg_inv::call(A);
  23487. }
  23488. auto [A_value, A_bdim] = unwrapTensorAtLevel(A, cur_level);
  23489. auto results = batch_rule(A_value, A_bdim);
  23490. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23491. }
  23492. template <typename batch_rule_t, batch_rule_t batch_rule>
  23493. at::Tensor inverse_generated_plumbing(const at::Tensor & self) {
  23494. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23495. auto maybe_layer = maybeCurrentDynamicLayer();
  23496. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23497. int64_t cur_level = maybe_layer->layerId();
  23498. if (!isBatchedAtLevel(self, cur_level)) {
  23499. return at::_ops::inverse::call(self);
  23500. }
  23501. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23502. auto results = batch_rule(self_value, self_bdim);
  23503. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23504. }
  23505. template <typename batch_rule_t, batch_rule_t batch_rule>
  23506. at::Tensor inner_generated_plumbing(const at::Tensor & self, const at::Tensor & other) {
  23507. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23508. auto maybe_layer = maybeCurrentDynamicLayer();
  23509. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23510. int64_t cur_level = maybe_layer->layerId();
  23511. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  23512. return at::_ops::inner::call(self, other);
  23513. }
  23514. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23515. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  23516. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim);
  23517. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23518. }
  23519. template <typename batch_rule_t, batch_rule_t batch_rule>
  23520. at::Tensor outer_generated_plumbing(const at::Tensor & self, const at::Tensor & vec2) {
  23521. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23522. auto maybe_layer = maybeCurrentDynamicLayer();
  23523. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23524. int64_t cur_level = maybe_layer->layerId();
  23525. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(vec2, cur_level)) {
  23526. return at::_ops::outer::call(self, vec2);
  23527. }
  23528. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23529. auto [vec2_value, vec2_bdim] = unwrapTensorAtLevel(vec2, cur_level);
  23530. auto results = batch_rule(self_value, self_bdim, vec2_value, vec2_bdim);
  23531. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23532. }
  23533. template <typename batch_rule_t, batch_rule_t batch_rule>
  23534. at::Tensor ger_generated_plumbing(const at::Tensor & self, const at::Tensor & vec2) {
  23535. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23536. auto maybe_layer = maybeCurrentDynamicLayer();
  23537. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23538. int64_t cur_level = maybe_layer->layerId();
  23539. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(vec2, cur_level)) {
  23540. return at::_ops::ger::call(self, vec2);
  23541. }
  23542. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23543. auto [vec2_value, vec2_bdim] = unwrapTensorAtLevel(vec2, cur_level);
  23544. auto results = batch_rule(self_value, self_bdim, vec2_value, vec2_bdim);
  23545. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23546. }
  23547. template <typename batch_rule_t, batch_rule_t batch_rule>
  23548. at::Tensor linalg_norm_generated_plumbing(const at::Tensor & self, const ::std::optional<at::Scalar> & ord, at::OptionalIntArrayRef dim, bool keepdim, ::std::optional<at::ScalarType> dtype) {
  23549. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23550. auto maybe_layer = maybeCurrentDynamicLayer();
  23551. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23552. int64_t cur_level = maybe_layer->layerId();
  23553. if (!isBatchedAtLevel(self, cur_level)) {
  23554. return at::_ops::linalg_norm::call(self, ord, dim, keepdim, dtype);
  23555. }
  23556. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23557. auto results = batch_rule(self_value, self_bdim, ord, dim, keepdim, dtype);
  23558. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23559. }
  23560. template <typename batch_rule_t, batch_rule_t batch_rule>
  23561. at::Tensor linalg_norm_ord_str_generated_plumbing(const at::Tensor & self, c10::string_view ord, at::OptionalIntArrayRef dim, bool keepdim, ::std::optional<at::ScalarType> dtype) {
  23562. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23563. auto maybe_layer = maybeCurrentDynamicLayer();
  23564. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23565. int64_t cur_level = maybe_layer->layerId();
  23566. if (!isBatchedAtLevel(self, cur_level)) {
  23567. return at::_ops::linalg_norm_ord_str::call(self, ord, dim, keepdim, dtype);
  23568. }
  23569. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23570. auto results = batch_rule(self_value, self_bdim, ord, dim, keepdim, dtype);
  23571. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23572. }
  23573. template <typename batch_rule_t, batch_rule_t batch_rule>
  23574. at::Tensor linalg_vector_norm_generated_plumbing(const at::Tensor & self, const at::Scalar & ord, at::OptionalIntArrayRef dim, bool keepdim, ::std::optional<at::ScalarType> dtype) {
  23575. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23576. auto maybe_layer = maybeCurrentDynamicLayer();
  23577. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23578. int64_t cur_level = maybe_layer->layerId();
  23579. if (!isBatchedAtLevel(self, cur_level)) {
  23580. return at::_ops::linalg_vector_norm::call(self, ord, dim, keepdim, dtype);
  23581. }
  23582. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23583. auto results = batch_rule(self_value, self_bdim, ord, dim, keepdim, dtype);
  23584. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23585. }
  23586. template <typename batch_rule_t, batch_rule_t batch_rule>
  23587. at::Tensor linalg_matrix_norm_generated_plumbing(const at::Tensor & self, const at::Scalar & ord, at::IntArrayRef dim, bool keepdim, ::std::optional<at::ScalarType> dtype) {
  23588. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23589. auto maybe_layer = maybeCurrentDynamicLayer();
  23590. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23591. int64_t cur_level = maybe_layer->layerId();
  23592. if (!isBatchedAtLevel(self, cur_level)) {
  23593. return at::_ops::linalg_matrix_norm::call(self, ord, dim, keepdim, dtype);
  23594. }
  23595. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23596. auto results = batch_rule(self_value, self_bdim, ord, dim, keepdim, dtype);
  23597. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23598. }
  23599. template <typename batch_rule_t, batch_rule_t batch_rule>
  23600. at::Tensor linalg_matrix_norm_str_ord_generated_plumbing(const at::Tensor & self, c10::string_view ord, at::IntArrayRef dim, bool keepdim, ::std::optional<at::ScalarType> dtype) {
  23601. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23602. auto maybe_layer = maybeCurrentDynamicLayer();
  23603. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23604. int64_t cur_level = maybe_layer->layerId();
  23605. if (!isBatchedAtLevel(self, cur_level)) {
  23606. return at::_ops::linalg_matrix_norm_str_ord::call(self, ord, dim, keepdim, dtype);
  23607. }
  23608. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23609. auto results = batch_rule(self_value, self_bdim, ord, dim, keepdim, dtype);
  23610. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23611. }
  23612. template <typename batch_rule_t, batch_rule_t batch_rule>
  23613. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> _linalg_svd_generated_plumbing(const at::Tensor & A, bool full_matrices, bool compute_uv, ::std::optional<c10::string_view> driver) {
  23614. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23615. auto maybe_layer = maybeCurrentDynamicLayer();
  23616. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23617. int64_t cur_level = maybe_layer->layerId();
  23618. if (!isBatchedAtLevel(A, cur_level)) {
  23619. return at::_ops::_linalg_svd::call(A, full_matrices, compute_uv, driver);
  23620. }
  23621. auto [A_value, A_bdim] = unwrapTensorAtLevel(A, cur_level);
  23622. auto results = batch_rule(A_value, A_bdim, full_matrices, compute_uv, driver);
  23623. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  23624. }
  23625. template <typename batch_rule_t, batch_rule_t batch_rule>
  23626. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> linalg_svd_generated_plumbing(const at::Tensor & A, bool full_matrices, ::std::optional<c10::string_view> driver) {
  23627. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23628. auto maybe_layer = maybeCurrentDynamicLayer();
  23629. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23630. int64_t cur_level = maybe_layer->layerId();
  23631. if (!isBatchedAtLevel(A, cur_level)) {
  23632. return at::_ops::linalg_svd::call(A, full_matrices, driver);
  23633. }
  23634. auto [A_value, A_bdim] = unwrapTensorAtLevel(A, cur_level);
  23635. auto results = batch_rule(A_value, A_bdim, full_matrices, driver);
  23636. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  23637. }
  23638. template <typename batch_rule_t, batch_rule_t batch_rule>
  23639. at::Tensor linalg_svdvals_generated_plumbing(const at::Tensor & A, ::std::optional<c10::string_view> driver) {
  23640. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23641. auto maybe_layer = maybeCurrentDynamicLayer();
  23642. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23643. int64_t cur_level = maybe_layer->layerId();
  23644. if (!isBatchedAtLevel(A, cur_level)) {
  23645. return at::_ops::linalg_svdvals::call(A, driver);
  23646. }
  23647. auto [A_value, A_bdim] = unwrapTensorAtLevel(A, cur_level);
  23648. auto results = batch_rule(A_value, A_bdim, driver);
  23649. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23650. }
  23651. template <typename batch_rule_t, batch_rule_t batch_rule>
  23652. at::Tensor linalg_cond_generated_plumbing(const at::Tensor & self, const ::std::optional<at::Scalar> & p) {
  23653. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23654. auto maybe_layer = maybeCurrentDynamicLayer();
  23655. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23656. int64_t cur_level = maybe_layer->layerId();
  23657. if (!isBatchedAtLevel(self, cur_level)) {
  23658. return at::_ops::linalg_cond::call(self, p);
  23659. }
  23660. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23661. auto results = batch_rule(self_value, self_bdim, p);
  23662. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23663. }
  23664. template <typename batch_rule_t, batch_rule_t batch_rule>
  23665. at::Tensor linalg_cond_p_str_generated_plumbing(const at::Tensor & self, c10::string_view p) {
  23666. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23667. auto maybe_layer = maybeCurrentDynamicLayer();
  23668. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23669. int64_t cur_level = maybe_layer->layerId();
  23670. if (!isBatchedAtLevel(self, cur_level)) {
  23671. return at::_ops::linalg_cond_p_str::call(self, p);
  23672. }
  23673. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23674. auto results = batch_rule(self_value, self_bdim, p);
  23675. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23676. }
  23677. template <typename batch_rule_t, batch_rule_t batch_rule>
  23678. at::Tensor linalg_pinv_atol_rtol_tensor_generated_plumbing(const at::Tensor & self, const ::std::optional<at::Tensor> & atol, const ::std::optional<at::Tensor> & rtol, bool hermitian) {
  23679. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23680. auto maybe_layer = maybeCurrentDynamicLayer();
  23681. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23682. int64_t cur_level = maybe_layer->layerId();
  23683. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(atol, cur_level) && !isBatchedAtLevel(rtol, cur_level)) {
  23684. return at::_ops::linalg_pinv_atol_rtol_tensor::call(self, atol, rtol, hermitian);
  23685. }
  23686. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23687. std::optional<Tensor> atol_value;
  23688. std::optional<int64_t> atol_bdim;
  23689. if (atol) {
  23690. std::tie(atol_value, atol_bdim) = unwrapTensorAtLevel(atol.value(), cur_level);
  23691. }
  23692. std::optional<Tensor> rtol_value;
  23693. std::optional<int64_t> rtol_bdim;
  23694. if (rtol) {
  23695. std::tie(rtol_value, rtol_bdim) = unwrapTensorAtLevel(rtol.value(), cur_level);
  23696. }
  23697. auto results = batch_rule(self_value, self_bdim, atol_value, atol_bdim, rtol_value, rtol_bdim, hermitian);
  23698. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23699. }
  23700. template <typename batch_rule_t, batch_rule_t batch_rule>
  23701. at::Tensor linalg_pinv_atol_rtol_float_generated_plumbing(const at::Tensor & self, ::std::optional<double> atol, ::std::optional<double> rtol, bool hermitian) {
  23702. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23703. auto maybe_layer = maybeCurrentDynamicLayer();
  23704. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23705. int64_t cur_level = maybe_layer->layerId();
  23706. if (!isBatchedAtLevel(self, cur_level)) {
  23707. return at::_ops::linalg_pinv_atol_rtol_float::call(self, atol, rtol, hermitian);
  23708. }
  23709. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23710. auto results = batch_rule(self_value, self_bdim, atol, rtol, hermitian);
  23711. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23712. }
  23713. template <typename batch_rule_t, batch_rule_t batch_rule>
  23714. at::Tensor linalg_pinv_generated_plumbing(const at::Tensor & self, double rcond, bool hermitian) {
  23715. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23716. auto maybe_layer = maybeCurrentDynamicLayer();
  23717. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23718. int64_t cur_level = maybe_layer->layerId();
  23719. if (!isBatchedAtLevel(self, cur_level)) {
  23720. return at::_ops::linalg_pinv::call(self, rcond, hermitian);
  23721. }
  23722. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23723. auto results = batch_rule(self_value, self_bdim, rcond, hermitian);
  23724. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23725. }
  23726. template <typename batch_rule_t, batch_rule_t batch_rule>
  23727. at::Tensor linalg_pinv_rcond_tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & rcond, bool hermitian) {
  23728. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23729. auto maybe_layer = maybeCurrentDynamicLayer();
  23730. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23731. int64_t cur_level = maybe_layer->layerId();
  23732. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(rcond, cur_level)) {
  23733. return at::_ops::linalg_pinv_rcond_tensor::call(self, rcond, hermitian);
  23734. }
  23735. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23736. auto [rcond_value, rcond_bdim] = unwrapTensorAtLevel(rcond, cur_level);
  23737. auto results = batch_rule(self_value, self_bdim, rcond_value, rcond_bdim, hermitian);
  23738. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23739. }
  23740. template <typename batch_rule_t, batch_rule_t batch_rule>
  23741. ::std::tuple<at::Tensor,at::Tensor,at::Tensor,at::Tensor> _linalg_solve_ex_generated_plumbing(const at::Tensor & A, const at::Tensor & B, bool left, bool check_errors) {
  23742. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23743. auto maybe_layer = maybeCurrentDynamicLayer();
  23744. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23745. int64_t cur_level = maybe_layer->layerId();
  23746. if (!isBatchedAtLevel(A, cur_level) && !isBatchedAtLevel(B, cur_level)) {
  23747. return at::_ops::_linalg_solve_ex::call(A, B, left, check_errors);
  23748. }
  23749. auto [A_value, A_bdim] = unwrapTensorAtLevel(A, cur_level);
  23750. auto [B_value, B_bdim] = unwrapTensorAtLevel(B, cur_level);
  23751. auto results = batch_rule(A_value, A_bdim, B_value, B_bdim, left, check_errors);
  23752. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level), makeBatched(std::get<6>(results), std::get<7>(results), cur_level));
  23753. }
  23754. template <typename batch_rule_t, batch_rule_t batch_rule>
  23755. ::std::tuple<at::Tensor,at::Tensor> linalg_solve_ex_generated_plumbing(const at::Tensor & A, const at::Tensor & B, bool left, bool check_errors) {
  23756. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23757. auto maybe_layer = maybeCurrentDynamicLayer();
  23758. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23759. int64_t cur_level = maybe_layer->layerId();
  23760. if (!isBatchedAtLevel(A, cur_level) && !isBatchedAtLevel(B, cur_level)) {
  23761. return at::_ops::linalg_solve_ex::call(A, B, left, check_errors);
  23762. }
  23763. auto [A_value, A_bdim] = unwrapTensorAtLevel(A, cur_level);
  23764. auto [B_value, B_bdim] = unwrapTensorAtLevel(B, cur_level);
  23765. auto results = batch_rule(A_value, A_bdim, B_value, B_bdim, left, check_errors);
  23766. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  23767. }
  23768. template <typename batch_rule_t, batch_rule_t batch_rule>
  23769. at::Tensor linalg_solve_generated_plumbing(const at::Tensor & A, const at::Tensor & B, bool left) {
  23770. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23771. auto maybe_layer = maybeCurrentDynamicLayer();
  23772. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23773. int64_t cur_level = maybe_layer->layerId();
  23774. if (!isBatchedAtLevel(A, cur_level) && !isBatchedAtLevel(B, cur_level)) {
  23775. return at::_ops::linalg_solve::call(A, B, left);
  23776. }
  23777. auto [A_value, A_bdim] = unwrapTensorAtLevel(A, cur_level);
  23778. auto [B_value, B_bdim] = unwrapTensorAtLevel(B, cur_level);
  23779. auto results = batch_rule(A_value, A_bdim, B_value, B_bdim, left);
  23780. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23781. }
  23782. template <typename batch_rule_t, batch_rule_t batch_rule>
  23783. at::Tensor _spsolve_generated_plumbing(const at::Tensor & A, const at::Tensor & B, bool left) {
  23784. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23785. auto maybe_layer = maybeCurrentDynamicLayer();
  23786. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23787. int64_t cur_level = maybe_layer->layerId();
  23788. if (!isBatchedAtLevel(A, cur_level) && !isBatchedAtLevel(B, cur_level)) {
  23789. return at::_ops::_spsolve::call(A, B, left);
  23790. }
  23791. auto [A_value, A_bdim] = unwrapTensorAtLevel(A, cur_level);
  23792. auto [B_value, B_bdim] = unwrapTensorAtLevel(B, cur_level);
  23793. auto results = batch_rule(A_value, A_bdim, B_value, B_bdim, left);
  23794. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23795. }
  23796. template <typename batch_rule_t, batch_rule_t batch_rule>
  23797. at::Tensor linalg_tensorinv_generated_plumbing(const at::Tensor & self, int64_t ind) {
  23798. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23799. auto maybe_layer = maybeCurrentDynamicLayer();
  23800. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23801. int64_t cur_level = maybe_layer->layerId();
  23802. if (!isBatchedAtLevel(self, cur_level)) {
  23803. return at::_ops::linalg_tensorinv::call(self, ind);
  23804. }
  23805. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23806. auto results = batch_rule(self_value, self_bdim, ind);
  23807. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23808. }
  23809. template <typename batch_rule_t, batch_rule_t batch_rule>
  23810. at::Tensor linalg_tensorsolve_generated_plumbing(const at::Tensor & self, const at::Tensor & other, at::OptionalIntArrayRef dims) {
  23811. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23812. auto maybe_layer = maybeCurrentDynamicLayer();
  23813. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23814. int64_t cur_level = maybe_layer->layerId();
  23815. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  23816. return at::_ops::linalg_tensorsolve::call(self, other, dims);
  23817. }
  23818. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23819. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  23820. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim, dims);
  23821. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23822. }
  23823. template <typename batch_rule_t, batch_rule_t batch_rule>
  23824. ::std::tuple<at::Tensor,at::Tensor> linalg_qr_generated_plumbing(const at::Tensor & A, c10::string_view mode) {
  23825. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23826. auto maybe_layer = maybeCurrentDynamicLayer();
  23827. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23828. int64_t cur_level = maybe_layer->layerId();
  23829. if (!isBatchedAtLevel(A, cur_level)) {
  23830. return at::_ops::linalg_qr::call(A, mode);
  23831. }
  23832. auto [A_value, A_bdim] = unwrapTensorAtLevel(A, cur_level);
  23833. auto results = batch_rule(A_value, A_bdim, mode);
  23834. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  23835. }
  23836. template <typename batch_rule_t, batch_rule_t batch_rule>
  23837. at::Tensor linalg_matrix_power_generated_plumbing(const at::Tensor & self, int64_t n) {
  23838. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23839. auto maybe_layer = maybeCurrentDynamicLayer();
  23840. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23841. int64_t cur_level = maybe_layer->layerId();
  23842. if (!isBatchedAtLevel(self, cur_level)) {
  23843. return at::_ops::linalg_matrix_power::call(self, n);
  23844. }
  23845. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23846. auto results = batch_rule(self_value, self_bdim, n);
  23847. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23848. }
  23849. template <typename batch_rule_t, batch_rule_t batch_rule>
  23850. at::Tensor linalg_matrix_rank_atol_rtol_tensor_generated_plumbing(const at::Tensor & input, const ::std::optional<at::Tensor> & atol, const ::std::optional<at::Tensor> & rtol, bool hermitian) {
  23851. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23852. auto maybe_layer = maybeCurrentDynamicLayer();
  23853. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23854. int64_t cur_level = maybe_layer->layerId();
  23855. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(atol, cur_level) && !isBatchedAtLevel(rtol, cur_level)) {
  23856. return at::_ops::linalg_matrix_rank_atol_rtol_tensor::call(input, atol, rtol, hermitian);
  23857. }
  23858. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  23859. std::optional<Tensor> atol_value;
  23860. std::optional<int64_t> atol_bdim;
  23861. if (atol) {
  23862. std::tie(atol_value, atol_bdim) = unwrapTensorAtLevel(atol.value(), cur_level);
  23863. }
  23864. std::optional<Tensor> rtol_value;
  23865. std::optional<int64_t> rtol_bdim;
  23866. if (rtol) {
  23867. std::tie(rtol_value, rtol_bdim) = unwrapTensorAtLevel(rtol.value(), cur_level);
  23868. }
  23869. auto results = batch_rule(input_value, input_bdim, atol_value, atol_bdim, rtol_value, rtol_bdim, hermitian);
  23870. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23871. }
  23872. template <typename batch_rule_t, batch_rule_t batch_rule>
  23873. at::Tensor linalg_matrix_rank_atol_rtol_float_generated_plumbing(const at::Tensor & self, ::std::optional<double> atol, ::std::optional<double> rtol, bool hermitian) {
  23874. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23875. auto maybe_layer = maybeCurrentDynamicLayer();
  23876. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23877. int64_t cur_level = maybe_layer->layerId();
  23878. if (!isBatchedAtLevel(self, cur_level)) {
  23879. return at::_ops::linalg_matrix_rank_atol_rtol_float::call(self, atol, rtol, hermitian);
  23880. }
  23881. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23882. auto results = batch_rule(self_value, self_bdim, atol, rtol, hermitian);
  23883. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23884. }
  23885. template <typename batch_rule_t, batch_rule_t batch_rule>
  23886. at::Tensor linalg_matrix_rank_generated_plumbing(const at::Tensor & self, double tol, bool hermitian) {
  23887. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23888. auto maybe_layer = maybeCurrentDynamicLayer();
  23889. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23890. int64_t cur_level = maybe_layer->layerId();
  23891. if (!isBatchedAtLevel(self, cur_level)) {
  23892. return at::_ops::linalg_matrix_rank::call(self, tol, hermitian);
  23893. }
  23894. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23895. auto results = batch_rule(self_value, self_bdim, tol, hermitian);
  23896. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23897. }
  23898. template <typename batch_rule_t, batch_rule_t batch_rule>
  23899. at::Tensor linalg_matrix_rank_tol_tensor_generated_plumbing(const at::Tensor & input, const at::Tensor & tol, bool hermitian) {
  23900. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23901. auto maybe_layer = maybeCurrentDynamicLayer();
  23902. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23903. int64_t cur_level = maybe_layer->layerId();
  23904. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(tol, cur_level)) {
  23905. return at::_ops::linalg_matrix_rank_tol_tensor::call(input, tol, hermitian);
  23906. }
  23907. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  23908. auto [tol_value, tol_bdim] = unwrapTensorAtLevel(tol, cur_level);
  23909. auto results = batch_rule(input_value, input_bdim, tol_value, tol_bdim, hermitian);
  23910. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23911. }
  23912. template <typename batch_rule_t, batch_rule_t batch_rule>
  23913. at::Tensor linalg_multi_dot_generated_plumbing(at::TensorList tensors) {
  23914. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23915. auto maybe_layer = maybeCurrentDynamicLayer();
  23916. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23917. int64_t cur_level = maybe_layer->layerId();
  23918. if (!isBatchedAtLevel(tensors, cur_level)) {
  23919. return at::_ops::linalg_multi_dot::call(tensors);
  23920. }
  23921. auto results = batch_rule(tensors);
  23922. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23923. }
  23924. template <typename batch_rule_t, batch_rule_t batch_rule>
  23925. at::Tensor nested_to_padded_tensor_generated_plumbing(const at::Tensor & self, double padding, at::OptionalIntArrayRef output_size) {
  23926. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23927. auto maybe_layer = maybeCurrentDynamicLayer();
  23928. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23929. int64_t cur_level = maybe_layer->layerId();
  23930. if (!isBatchedAtLevel(self, cur_level)) {
  23931. return at::_ops::nested_to_padded_tensor::call(self, padding, output_size);
  23932. }
  23933. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23934. auto results = batch_rule(self_value, self_bdim, padding, output_size);
  23935. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23936. }
  23937. template <typename batch_rule_t, batch_rule_t batch_rule>
  23938. at::Tensor _test_serialization_subcmul_generated_plumbing(const at::Tensor & self, const at::Tensor & other, const at::Scalar & alpha) {
  23939. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23940. auto maybe_layer = maybeCurrentDynamicLayer();
  23941. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23942. int64_t cur_level = maybe_layer->layerId();
  23943. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level)) {
  23944. return at::_ops::_test_serialization_subcmul::call(self, other, alpha);
  23945. }
  23946. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23947. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  23948. auto results = batch_rule(self_value, self_bdim, other_value, other_bdim, alpha);
  23949. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23950. }
  23951. template <typename batch_rule_t, batch_rule_t batch_rule>
  23952. at::Tensor _test_parallel_materialize_generated_plumbing(const at::Tensor & self, int64_t num_parallel, bool skip_first) {
  23953. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23954. auto maybe_layer = maybeCurrentDynamicLayer();
  23955. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23956. int64_t cur_level = maybe_layer->layerId();
  23957. if (!isBatchedAtLevel(self, cur_level)) {
  23958. return at::_ops::_test_parallel_materialize::call(self, num_parallel, skip_first);
  23959. }
  23960. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  23961. auto results = batch_rule(self_value, self_bdim, num_parallel, skip_first);
  23962. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23963. }
  23964. template <typename batch_rule_t, batch_rule_t batch_rule>
  23965. at::Tensor _test_optional_intlist_generated_plumbing(const at::Tensor & values, at::OptionalIntArrayRef addends) {
  23966. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23967. auto maybe_layer = maybeCurrentDynamicLayer();
  23968. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23969. int64_t cur_level = maybe_layer->layerId();
  23970. if (!isBatchedAtLevel(values, cur_level)) {
  23971. return at::_ops::_test_optional_intlist::call(values, addends);
  23972. }
  23973. auto [values_value, values_bdim] = unwrapTensorAtLevel(values, cur_level);
  23974. auto results = batch_rule(values_value, values_bdim, addends);
  23975. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23976. }
  23977. template <typename batch_rule_t, batch_rule_t batch_rule>
  23978. at::Tensor _test_optional_filled_intlist_generated_plumbing(const at::Tensor & values, at::OptionalIntArrayRef addends) {
  23979. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23980. auto maybe_layer = maybeCurrentDynamicLayer();
  23981. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23982. int64_t cur_level = maybe_layer->layerId();
  23983. if (!isBatchedAtLevel(values, cur_level)) {
  23984. return at::_ops::_test_optional_filled_intlist::call(values, addends);
  23985. }
  23986. auto [values_value, values_bdim] = unwrapTensorAtLevel(values, cur_level);
  23987. auto results = batch_rule(values_value, values_bdim, addends);
  23988. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  23989. }
  23990. template <typename batch_rule_t, batch_rule_t batch_rule>
  23991. at::Tensor _test_optional_floatlist_generated_plumbing(const at::Tensor & values, ::std::optional<at::ArrayRef<double>> addends) {
  23992. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  23993. auto maybe_layer = maybeCurrentDynamicLayer();
  23994. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  23995. int64_t cur_level = maybe_layer->layerId();
  23996. if (!isBatchedAtLevel(values, cur_level)) {
  23997. return at::_ops::_test_optional_floatlist::call(values, addends);
  23998. }
  23999. auto [values_value, values_bdim] = unwrapTensorAtLevel(values, cur_level);
  24000. auto results = batch_rule(values_value, values_bdim, addends);
  24001. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24002. }
  24003. template <typename batch_rule_t, batch_rule_t batch_rule>
  24004. at::Tensor _test_string_default_generated_plumbing(const at::Tensor & dummy, c10::string_view a, c10::string_view b) {
  24005. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24006. auto maybe_layer = maybeCurrentDynamicLayer();
  24007. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24008. int64_t cur_level = maybe_layer->layerId();
  24009. if (!isBatchedAtLevel(dummy, cur_level)) {
  24010. return at::_ops::_test_string_default::call(dummy, a, b);
  24011. }
  24012. auto [dummy_value, dummy_bdim] = unwrapTensorAtLevel(dummy, cur_level);
  24013. auto results = batch_rule(dummy_value, dummy_bdim, a, b);
  24014. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24015. }
  24016. template <typename batch_rule_t, batch_rule_t batch_rule>
  24017. at::Tensor _test_ambiguous_defaults_a_generated_plumbing(const at::Tensor & dummy, int64_t a, int64_t b) {
  24018. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24019. auto maybe_layer = maybeCurrentDynamicLayer();
  24020. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24021. int64_t cur_level = maybe_layer->layerId();
  24022. if (!isBatchedAtLevel(dummy, cur_level)) {
  24023. return at::_ops::_test_ambiguous_defaults_a::call(dummy, a, b);
  24024. }
  24025. auto [dummy_value, dummy_bdim] = unwrapTensorAtLevel(dummy, cur_level);
  24026. auto results = batch_rule(dummy_value, dummy_bdim, a, b);
  24027. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24028. }
  24029. template <typename batch_rule_t, batch_rule_t batch_rule>
  24030. at::Tensor _test_ambiguous_defaults_b_generated_plumbing(const at::Tensor & dummy, int64_t a, c10::string_view b) {
  24031. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24032. auto maybe_layer = maybeCurrentDynamicLayer();
  24033. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24034. int64_t cur_level = maybe_layer->layerId();
  24035. if (!isBatchedAtLevel(dummy, cur_level)) {
  24036. return at::_ops::_test_ambiguous_defaults_b::call(dummy, a, b);
  24037. }
  24038. auto [dummy_value, dummy_bdim] = unwrapTensorAtLevel(dummy, cur_level);
  24039. auto results = batch_rule(dummy_value, dummy_bdim, a, b);
  24040. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24041. }
  24042. template <typename batch_rule_t, batch_rule_t batch_rule>
  24043. at::Tensor _test_warn_in_autograd_generated_plumbing(const at::Tensor & self) {
  24044. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24045. auto maybe_layer = maybeCurrentDynamicLayer();
  24046. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24047. int64_t cur_level = maybe_layer->layerId();
  24048. if (!isBatchedAtLevel(self, cur_level)) {
  24049. return at::_ops::_test_warn_in_autograd::call(self);
  24050. }
  24051. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24052. auto results = batch_rule(self_value, self_bdim);
  24053. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24054. }
  24055. template <typename batch_rule_t, batch_rule_t batch_rule>
  24056. at::Tensor _test_autograd_multiple_dispatch_fullcoverage_generated_plumbing(const at::Tensor & self) {
  24057. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24058. auto maybe_layer = maybeCurrentDynamicLayer();
  24059. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24060. int64_t cur_level = maybe_layer->layerId();
  24061. if (!isBatchedAtLevel(self, cur_level)) {
  24062. return at::_ops::_test_autograd_multiple_dispatch_fullcoverage::call(self);
  24063. }
  24064. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24065. auto results = batch_rule(self_value, self_bdim);
  24066. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24067. }
  24068. template <typename batch_rule_t, batch_rule_t batch_rule>
  24069. at::Tensor _test_autograd_multiple_dispatch_ntonly_generated_plumbing(const at::Tensor & self, bool b) {
  24070. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24071. auto maybe_layer = maybeCurrentDynamicLayer();
  24072. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24073. int64_t cur_level = maybe_layer->layerId();
  24074. if (!isBatchedAtLevel(self, cur_level)) {
  24075. return at::_ops::_test_autograd_multiple_dispatch_ntonly::call(self, b);
  24076. }
  24077. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24078. auto results = batch_rule(self_value, self_bdim, b);
  24079. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24080. }
  24081. template <typename batch_rule_t, batch_rule_t batch_rule>
  24082. at::Tensor _test_autograd_multiple_dispatch_view_generated_plumbing(const at::Tensor & self) {
  24083. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24084. auto maybe_layer = maybeCurrentDynamicLayer();
  24085. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24086. int64_t cur_level = maybe_layer->layerId();
  24087. if (!isBatchedAtLevel(self, cur_level)) {
  24088. return at::_ops::_test_autograd_multiple_dispatch_view::call(self);
  24089. }
  24090. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24091. auto results = batch_rule(self_value, self_bdim);
  24092. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24093. }
  24094. template <typename batch_rule_t, batch_rule_t batch_rule>
  24095. at::Tensor _test_autograd_multiple_dispatch_view_copy_generated_plumbing(const at::Tensor & self) {
  24096. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24097. auto maybe_layer = maybeCurrentDynamicLayer();
  24098. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24099. int64_t cur_level = maybe_layer->layerId();
  24100. if (!isBatchedAtLevel(self, cur_level)) {
  24101. return at::_ops::_test_autograd_multiple_dispatch_view_copy::call(self);
  24102. }
  24103. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24104. auto results = batch_rule(self_value, self_bdim);
  24105. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24106. }
  24107. template <typename batch_rule_t, batch_rule_t batch_rule>
  24108. at::Tensor segment_reduce_generated_plumbing(const at::Tensor & data, c10::string_view reduce, const ::std::optional<at::Tensor> & lengths, const ::std::optional<at::Tensor> & indices, const ::std::optional<at::Tensor> & offsets, int64_t axis, bool unsafe, const ::std::optional<at::Scalar> & initial) {
  24109. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24110. auto maybe_layer = maybeCurrentDynamicLayer();
  24111. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24112. int64_t cur_level = maybe_layer->layerId();
  24113. if (!isBatchedAtLevel(data, cur_level) && !isBatchedAtLevel(lengths, cur_level) && !isBatchedAtLevel(indices, cur_level) && !isBatchedAtLevel(offsets, cur_level)) {
  24114. return at::_ops::segment_reduce::call(data, reduce, lengths, indices, offsets, axis, unsafe, initial);
  24115. }
  24116. auto [data_value, data_bdim] = unwrapTensorAtLevel(data, cur_level);
  24117. std::optional<Tensor> lengths_value;
  24118. std::optional<int64_t> lengths_bdim;
  24119. if (lengths) {
  24120. std::tie(lengths_value, lengths_bdim) = unwrapTensorAtLevel(lengths.value(), cur_level);
  24121. }
  24122. std::optional<Tensor> indices_value;
  24123. std::optional<int64_t> indices_bdim;
  24124. if (indices) {
  24125. std::tie(indices_value, indices_bdim) = unwrapTensorAtLevel(indices.value(), cur_level);
  24126. }
  24127. std::optional<Tensor> offsets_value;
  24128. std::optional<int64_t> offsets_bdim;
  24129. if (offsets) {
  24130. std::tie(offsets_value, offsets_bdim) = unwrapTensorAtLevel(offsets.value(), cur_level);
  24131. }
  24132. auto results = batch_rule(data_value, data_bdim, reduce, lengths_value, lengths_bdim, indices_value, indices_bdim, offsets_value, offsets_bdim, axis, unsafe, initial);
  24133. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24134. }
  24135. template <typename batch_rule_t, batch_rule_t batch_rule>
  24136. at::Tensor _segment_reduce_backward_generated_plumbing(const at::Tensor & grad, const at::Tensor & output, const at::Tensor & data, c10::string_view reduce, const ::std::optional<at::Tensor> & lengths, const ::std::optional<at::Tensor> & offsets, int64_t axis, const ::std::optional<at::Scalar> & initial) {
  24137. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24138. auto maybe_layer = maybeCurrentDynamicLayer();
  24139. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24140. int64_t cur_level = maybe_layer->layerId();
  24141. if (!isBatchedAtLevel(grad, cur_level) && !isBatchedAtLevel(output, cur_level) && !isBatchedAtLevel(data, cur_level) && !isBatchedAtLevel(lengths, cur_level) && !isBatchedAtLevel(offsets, cur_level)) {
  24142. return at::_ops::_segment_reduce_backward::call(grad, output, data, reduce, lengths, offsets, axis, initial);
  24143. }
  24144. auto [grad_value, grad_bdim] = unwrapTensorAtLevel(grad, cur_level);
  24145. auto [output_value, output_bdim] = unwrapTensorAtLevel(output, cur_level);
  24146. auto [data_value, data_bdim] = unwrapTensorAtLevel(data, cur_level);
  24147. std::optional<Tensor> lengths_value;
  24148. std::optional<int64_t> lengths_bdim;
  24149. if (lengths) {
  24150. std::tie(lengths_value, lengths_bdim) = unwrapTensorAtLevel(lengths.value(), cur_level);
  24151. }
  24152. std::optional<Tensor> offsets_value;
  24153. std::optional<int64_t> offsets_bdim;
  24154. if (offsets) {
  24155. std::tie(offsets_value, offsets_bdim) = unwrapTensorAtLevel(offsets.value(), cur_level);
  24156. }
  24157. auto results = batch_rule(grad_value, grad_bdim, output_value, output_bdim, data_value, data_bdim, reduce, lengths_value, lengths_bdim, offsets_value, offsets_bdim, axis, initial);
  24158. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24159. }
  24160. template <typename batch_rule_t, batch_rule_t batch_rule>
  24161. at::Tensor pad_sequence_generated_plumbing(at::TensorList sequences, bool batch_first, double padding_value, c10::string_view padding_side) {
  24162. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24163. auto maybe_layer = maybeCurrentDynamicLayer();
  24164. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24165. int64_t cur_level = maybe_layer->layerId();
  24166. if (!isBatchedAtLevel(sequences, cur_level)) {
  24167. return at::_ops::pad_sequence::call(sequences, batch_first, padding_value, padding_side);
  24168. }
  24169. auto results = batch_rule(sequences, batch_first, padding_value, padding_side);
  24170. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24171. }
  24172. template <typename batch_rule_t, batch_rule_t batch_rule>
  24173. at::Tensor flatten_dense_tensors_generated_plumbing(at::TensorList tensors) {
  24174. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24175. auto maybe_layer = maybeCurrentDynamicLayer();
  24176. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24177. int64_t cur_level = maybe_layer->layerId();
  24178. if (!isBatchedAtLevel(tensors, cur_level)) {
  24179. return at::_ops::flatten_dense_tensors::call(tensors);
  24180. }
  24181. auto results = batch_rule(tensors);
  24182. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24183. }
  24184. template <typename batch_rule_t, batch_rule_t batch_rule>
  24185. ::std::vector<at::Tensor> unflatten_dense_tensors_generated_plumbing(const at::Tensor & flat, at::TensorList tensors) {
  24186. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24187. auto maybe_layer = maybeCurrentDynamicLayer();
  24188. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24189. int64_t cur_level = maybe_layer->layerId();
  24190. if (!isBatchedAtLevel(flat, cur_level) && !isBatchedAtLevel(tensors, cur_level)) {
  24191. return at::_ops::unflatten_dense_tensors::call(flat, tensors);
  24192. }
  24193. auto [flat_value, flat_bdim] = unwrapTensorAtLevel(flat, cur_level);
  24194. auto results = batch_rule(flat_value, flat_bdim, tensors);
  24195. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  24196. }
  24197. template <typename batch_rule_t, batch_rule_t batch_rule>
  24198. at::Tensor _nested_tensor_from_tensor_list_generated_plumbing(at::TensorList list, ::std::optional<at::ScalarType> dtype, ::std::optional<at::Layout> layout, ::std::optional<at::Device> device, ::std::optional<bool> pin_memory) {
  24199. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24200. auto maybe_layer = maybeCurrentDynamicLayer();
  24201. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24202. int64_t cur_level = maybe_layer->layerId();
  24203. if (!isBatchedAtLevel(list, cur_level)) {
  24204. return at::_ops::_nested_tensor_from_tensor_list::call(list, dtype, layout, device, pin_memory);
  24205. }
  24206. auto results = batch_rule(list, dtype, layout, device, pin_memory);
  24207. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24208. }
  24209. template <typename batch_rule_t, batch_rule_t batch_rule>
  24210. at::Tensor _fw_primal_copy_generated_plumbing(const at::Tensor & self, int64_t level) {
  24211. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24212. auto maybe_layer = maybeCurrentDynamicLayer();
  24213. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24214. int64_t cur_level = maybe_layer->layerId();
  24215. if (!isBatchedAtLevel(self, cur_level)) {
  24216. return at::_ops::_fw_primal_copy::call(self, level);
  24217. }
  24218. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24219. auto results = batch_rule(self_value, self_bdim, level);
  24220. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24221. }
  24222. template <typename batch_rule_t, batch_rule_t batch_rule>
  24223. at::Tensor _make_dual_copy_generated_plumbing(const at::Tensor & primal, const at::Tensor & tangent, int64_t level) {
  24224. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24225. auto maybe_layer = maybeCurrentDynamicLayer();
  24226. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24227. int64_t cur_level = maybe_layer->layerId();
  24228. if (!isBatchedAtLevel(primal, cur_level) && !isBatchedAtLevel(tangent, cur_level)) {
  24229. return at::_ops::_make_dual_copy::call(primal, tangent, level);
  24230. }
  24231. auto [primal_value, primal_bdim] = unwrapTensorAtLevel(primal, cur_level);
  24232. auto [tangent_value, tangent_bdim] = unwrapTensorAtLevel(tangent, cur_level);
  24233. auto results = batch_rule(primal_value, primal_bdim, tangent_value, tangent_bdim, level);
  24234. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24235. }
  24236. template <typename batch_rule_t, batch_rule_t batch_rule>
  24237. at::Tensor view_as_real_copy_generated_plumbing(const at::Tensor & self) {
  24238. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24239. auto maybe_layer = maybeCurrentDynamicLayer();
  24240. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24241. int64_t cur_level = maybe_layer->layerId();
  24242. if (!isBatchedAtLevel(self, cur_level)) {
  24243. return at::_ops::view_as_real_copy::call(self);
  24244. }
  24245. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24246. auto results = batch_rule(self_value, self_bdim);
  24247. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24248. }
  24249. template <typename batch_rule_t, batch_rule_t batch_rule>
  24250. at::Tensor view_as_complex_copy_generated_plumbing(const at::Tensor & self) {
  24251. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24252. auto maybe_layer = maybeCurrentDynamicLayer();
  24253. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24254. int64_t cur_level = maybe_layer->layerId();
  24255. if (!isBatchedAtLevel(self, cur_level)) {
  24256. return at::_ops::view_as_complex_copy::call(self);
  24257. }
  24258. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24259. auto results = batch_rule(self_value, self_bdim);
  24260. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24261. }
  24262. template <typename batch_rule_t, batch_rule_t batch_rule>
  24263. at::Tensor _conj_copy_generated_plumbing(const at::Tensor & self) {
  24264. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24265. auto maybe_layer = maybeCurrentDynamicLayer();
  24266. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24267. int64_t cur_level = maybe_layer->layerId();
  24268. if (!isBatchedAtLevel(self, cur_level)) {
  24269. return at::_ops::_conj_copy::call(self);
  24270. }
  24271. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24272. auto results = batch_rule(self_value, self_bdim);
  24273. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24274. }
  24275. template <typename batch_rule_t, batch_rule_t batch_rule>
  24276. at::Tensor _neg_view_copy_generated_plumbing(const at::Tensor & self) {
  24277. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24278. auto maybe_layer = maybeCurrentDynamicLayer();
  24279. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24280. int64_t cur_level = maybe_layer->layerId();
  24281. if (!isBatchedAtLevel(self, cur_level)) {
  24282. return at::_ops::_neg_view_copy::call(self);
  24283. }
  24284. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24285. auto results = batch_rule(self_value, self_bdim);
  24286. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24287. }
  24288. template <typename batch_rule_t, batch_rule_t batch_rule>
  24289. at::Tensor as_strided_copy_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef size, c10::SymIntArrayRef stride, ::std::optional<c10::SymInt> storage_offset) {
  24290. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24291. auto maybe_layer = maybeCurrentDynamicLayer();
  24292. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24293. int64_t cur_level = maybe_layer->layerId();
  24294. if (!isBatchedAtLevel(self, cur_level)) {
  24295. return at::_ops::as_strided_copy::call(self, size, stride, storage_offset);
  24296. }
  24297. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24298. auto results = batch_rule(self_value, self_bdim, size, stride, storage_offset);
  24299. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24300. }
  24301. template <typename batch_rule_t, batch_rule_t batch_rule>
  24302. at::Tensor _sparse_broadcast_to_copy_generated_plumbing(const at::Tensor & self, at::IntArrayRef size) {
  24303. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24304. auto maybe_layer = maybeCurrentDynamicLayer();
  24305. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24306. int64_t cur_level = maybe_layer->layerId();
  24307. if (!isBatchedAtLevel(self, cur_level)) {
  24308. return at::_ops::_sparse_broadcast_to_copy::call(self, size);
  24309. }
  24310. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24311. auto results = batch_rule(self_value, self_bdim, size);
  24312. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24313. }
  24314. template <typename batch_rule_t, batch_rule_t batch_rule>
  24315. at::Tensor diagonal_copy_generated_plumbing(const at::Tensor & self, int64_t offset, int64_t dim1, int64_t dim2) {
  24316. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24317. auto maybe_layer = maybeCurrentDynamicLayer();
  24318. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24319. int64_t cur_level = maybe_layer->layerId();
  24320. if (!isBatchedAtLevel(self, cur_level)) {
  24321. return at::_ops::diagonal_copy::call(self, offset, dim1, dim2);
  24322. }
  24323. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24324. auto results = batch_rule(self_value, self_bdim, offset, dim1, dim2);
  24325. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24326. }
  24327. template <typename batch_rule_t, batch_rule_t batch_rule>
  24328. at::Tensor expand_copy_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef size, bool implicit) {
  24329. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24330. auto maybe_layer = maybeCurrentDynamicLayer();
  24331. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24332. int64_t cur_level = maybe_layer->layerId();
  24333. if (!isBatchedAtLevel(self, cur_level)) {
  24334. return at::_ops::expand_copy::call(self, size, implicit);
  24335. }
  24336. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24337. auto results = batch_rule(self_value, self_bdim, size, implicit);
  24338. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24339. }
  24340. template <typename batch_rule_t, batch_rule_t batch_rule>
  24341. at::Tensor permute_copy_generated_plumbing(const at::Tensor & self, at::IntArrayRef dims) {
  24342. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24343. auto maybe_layer = maybeCurrentDynamicLayer();
  24344. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24345. int64_t cur_level = maybe_layer->layerId();
  24346. if (!isBatchedAtLevel(self, cur_level)) {
  24347. return at::_ops::permute_copy::call(self, dims);
  24348. }
  24349. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24350. auto results = batch_rule(self_value, self_bdim, dims);
  24351. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24352. }
  24353. template <typename batch_rule_t, batch_rule_t batch_rule>
  24354. at::Tensor _reshape_alias_copy_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef size, c10::SymIntArrayRef stride) {
  24355. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24356. auto maybe_layer = maybeCurrentDynamicLayer();
  24357. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24358. int64_t cur_level = maybe_layer->layerId();
  24359. if (!isBatchedAtLevel(self, cur_level)) {
  24360. return at::_ops::_reshape_alias_copy::call(self, size, stride);
  24361. }
  24362. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24363. auto results = batch_rule(self_value, self_bdim, size, stride);
  24364. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24365. }
  24366. template <typename batch_rule_t, batch_rule_t batch_rule>
  24367. at::Tensor select_copy_int_generated_plumbing(const at::Tensor & self, int64_t dim, c10::SymInt index) {
  24368. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24369. auto maybe_layer = maybeCurrentDynamicLayer();
  24370. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24371. int64_t cur_level = maybe_layer->layerId();
  24372. if (!isBatchedAtLevel(self, cur_level)) {
  24373. return at::_ops::select_copy_int::call(self, dim, index);
  24374. }
  24375. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24376. auto results = batch_rule(self_value, self_bdim, dim, index);
  24377. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24378. }
  24379. template <typename batch_rule_t, batch_rule_t batch_rule>
  24380. at::Tensor detach_copy_generated_plumbing(const at::Tensor & self) {
  24381. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24382. auto maybe_layer = maybeCurrentDynamicLayer();
  24383. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24384. int64_t cur_level = maybe_layer->layerId();
  24385. if (!isBatchedAtLevel(self, cur_level)) {
  24386. return at::_ops::detach_copy::call(self);
  24387. }
  24388. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24389. auto results = batch_rule(self_value, self_bdim);
  24390. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24391. }
  24392. template <typename batch_rule_t, batch_rule_t batch_rule>
  24393. at::Tensor slice_copy_Tensor_generated_plumbing(const at::Tensor & self, int64_t dim, ::std::optional<c10::SymInt> start, ::std::optional<c10::SymInt> end, c10::SymInt step) {
  24394. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24395. auto maybe_layer = maybeCurrentDynamicLayer();
  24396. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24397. int64_t cur_level = maybe_layer->layerId();
  24398. if (!isBatchedAtLevel(self, cur_level)) {
  24399. return at::_ops::slice_copy_Tensor::call(self, dim, start, end, step);
  24400. }
  24401. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24402. auto results = batch_rule(self_value, self_bdim, dim, start, end, step);
  24403. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24404. }
  24405. template <typename batch_rule_t, batch_rule_t batch_rule>
  24406. ::std::vector<at::Tensor> split_copy_Tensor_generated_plumbing(const at::Tensor & self, c10::SymInt split_size, int64_t dim) {
  24407. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24408. auto maybe_layer = maybeCurrentDynamicLayer();
  24409. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24410. int64_t cur_level = maybe_layer->layerId();
  24411. if (!isBatchedAtLevel(self, cur_level)) {
  24412. return at::_ops::split_copy_Tensor::call(self, split_size, dim);
  24413. }
  24414. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24415. auto results = batch_rule(self_value, self_bdim, split_size, dim);
  24416. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  24417. }
  24418. template <typename batch_rule_t, batch_rule_t batch_rule>
  24419. ::std::vector<at::Tensor> split_with_sizes_copy_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef split_sizes, int64_t dim) {
  24420. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24421. auto maybe_layer = maybeCurrentDynamicLayer();
  24422. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24423. int64_t cur_level = maybe_layer->layerId();
  24424. if (!isBatchedAtLevel(self, cur_level)) {
  24425. return at::_ops::split_with_sizes_copy::call(self, split_sizes, dim);
  24426. }
  24427. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24428. auto results = batch_rule(self_value, self_bdim, split_sizes, dim);
  24429. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  24430. }
  24431. template <typename batch_rule_t, batch_rule_t batch_rule>
  24432. at::Tensor squeeze_copy_generated_plumbing(const at::Tensor & self) {
  24433. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24434. auto maybe_layer = maybeCurrentDynamicLayer();
  24435. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24436. int64_t cur_level = maybe_layer->layerId();
  24437. if (!isBatchedAtLevel(self, cur_level)) {
  24438. return at::_ops::squeeze_copy::call(self);
  24439. }
  24440. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24441. auto results = batch_rule(self_value, self_bdim);
  24442. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24443. }
  24444. template <typename batch_rule_t, batch_rule_t batch_rule>
  24445. at::Tensor squeeze_copy_dim_generated_plumbing(const at::Tensor & self, int64_t dim) {
  24446. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24447. auto maybe_layer = maybeCurrentDynamicLayer();
  24448. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24449. int64_t cur_level = maybe_layer->layerId();
  24450. if (!isBatchedAtLevel(self, cur_level)) {
  24451. return at::_ops::squeeze_copy_dim::call(self, dim);
  24452. }
  24453. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24454. auto results = batch_rule(self_value, self_bdim, dim);
  24455. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24456. }
  24457. template <typename batch_rule_t, batch_rule_t batch_rule>
  24458. at::Tensor squeeze_copy_dims_generated_plumbing(const at::Tensor & self, at::IntArrayRef dim) {
  24459. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24460. auto maybe_layer = maybeCurrentDynamicLayer();
  24461. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24462. int64_t cur_level = maybe_layer->layerId();
  24463. if (!isBatchedAtLevel(self, cur_level)) {
  24464. return at::_ops::squeeze_copy_dims::call(self, dim);
  24465. }
  24466. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24467. auto results = batch_rule(self_value, self_bdim, dim);
  24468. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24469. }
  24470. template <typename batch_rule_t, batch_rule_t batch_rule>
  24471. at::Tensor t_copy_generated_plumbing(const at::Tensor & self) {
  24472. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24473. auto maybe_layer = maybeCurrentDynamicLayer();
  24474. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24475. int64_t cur_level = maybe_layer->layerId();
  24476. if (!isBatchedAtLevel(self, cur_level)) {
  24477. return at::_ops::t_copy::call(self);
  24478. }
  24479. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24480. auto results = batch_rule(self_value, self_bdim);
  24481. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24482. }
  24483. template <typename batch_rule_t, batch_rule_t batch_rule>
  24484. at::Tensor transpose_copy_int_generated_plumbing(const at::Tensor & self, int64_t dim0, int64_t dim1) {
  24485. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24486. auto maybe_layer = maybeCurrentDynamicLayer();
  24487. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24488. int64_t cur_level = maybe_layer->layerId();
  24489. if (!isBatchedAtLevel(self, cur_level)) {
  24490. return at::_ops::transpose_copy_int::call(self, dim0, dim1);
  24491. }
  24492. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24493. auto results = batch_rule(self_value, self_bdim, dim0, dim1);
  24494. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24495. }
  24496. template <typename batch_rule_t, batch_rule_t batch_rule>
  24497. at::Tensor unsqueeze_copy_generated_plumbing(const at::Tensor & self, int64_t dim) {
  24498. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24499. auto maybe_layer = maybeCurrentDynamicLayer();
  24500. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24501. int64_t cur_level = maybe_layer->layerId();
  24502. if (!isBatchedAtLevel(self, cur_level)) {
  24503. return at::_ops::unsqueeze_copy::call(self, dim);
  24504. }
  24505. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24506. auto results = batch_rule(self_value, self_bdim, dim);
  24507. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24508. }
  24509. template <typename batch_rule_t, batch_rule_t batch_rule>
  24510. at::Tensor _indices_copy_generated_plumbing(const at::Tensor & self) {
  24511. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24512. auto maybe_layer = maybeCurrentDynamicLayer();
  24513. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24514. int64_t cur_level = maybe_layer->layerId();
  24515. if (!isBatchedAtLevel(self, cur_level)) {
  24516. return at::_ops::_indices_copy::call(self);
  24517. }
  24518. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24519. auto results = batch_rule(self_value, self_bdim);
  24520. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24521. }
  24522. template <typename batch_rule_t, batch_rule_t batch_rule>
  24523. at::Tensor _values_copy_generated_plumbing(const at::Tensor & self) {
  24524. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24525. auto maybe_layer = maybeCurrentDynamicLayer();
  24526. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24527. int64_t cur_level = maybe_layer->layerId();
  24528. if (!isBatchedAtLevel(self, cur_level)) {
  24529. return at::_ops::_values_copy::call(self);
  24530. }
  24531. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24532. auto results = batch_rule(self_value, self_bdim);
  24533. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24534. }
  24535. template <typename batch_rule_t, batch_rule_t batch_rule>
  24536. at::Tensor indices_copy_generated_plumbing(const at::Tensor & self) {
  24537. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24538. auto maybe_layer = maybeCurrentDynamicLayer();
  24539. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24540. int64_t cur_level = maybe_layer->layerId();
  24541. if (!isBatchedAtLevel(self, cur_level)) {
  24542. return at::_ops::indices_copy::call(self);
  24543. }
  24544. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24545. auto results = batch_rule(self_value, self_bdim);
  24546. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24547. }
  24548. template <typename batch_rule_t, batch_rule_t batch_rule>
  24549. at::Tensor values_copy_generated_plumbing(const at::Tensor & self) {
  24550. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24551. auto maybe_layer = maybeCurrentDynamicLayer();
  24552. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24553. int64_t cur_level = maybe_layer->layerId();
  24554. if (!isBatchedAtLevel(self, cur_level)) {
  24555. return at::_ops::values_copy::call(self);
  24556. }
  24557. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24558. auto results = batch_rule(self_value, self_bdim);
  24559. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24560. }
  24561. template <typename batch_rule_t, batch_rule_t batch_rule>
  24562. at::Tensor crow_indices_copy_generated_plumbing(const at::Tensor & self) {
  24563. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24564. auto maybe_layer = maybeCurrentDynamicLayer();
  24565. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24566. int64_t cur_level = maybe_layer->layerId();
  24567. if (!isBatchedAtLevel(self, cur_level)) {
  24568. return at::_ops::crow_indices_copy::call(self);
  24569. }
  24570. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24571. auto results = batch_rule(self_value, self_bdim);
  24572. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24573. }
  24574. template <typename batch_rule_t, batch_rule_t batch_rule>
  24575. at::Tensor col_indices_copy_generated_plumbing(const at::Tensor & self) {
  24576. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24577. auto maybe_layer = maybeCurrentDynamicLayer();
  24578. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24579. int64_t cur_level = maybe_layer->layerId();
  24580. if (!isBatchedAtLevel(self, cur_level)) {
  24581. return at::_ops::col_indices_copy::call(self);
  24582. }
  24583. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24584. auto results = batch_rule(self_value, self_bdim);
  24585. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24586. }
  24587. template <typename batch_rule_t, batch_rule_t batch_rule>
  24588. at::Tensor ccol_indices_copy_generated_plumbing(const at::Tensor & self) {
  24589. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24590. auto maybe_layer = maybeCurrentDynamicLayer();
  24591. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24592. int64_t cur_level = maybe_layer->layerId();
  24593. if (!isBatchedAtLevel(self, cur_level)) {
  24594. return at::_ops::ccol_indices_copy::call(self);
  24595. }
  24596. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24597. auto results = batch_rule(self_value, self_bdim);
  24598. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24599. }
  24600. template <typename batch_rule_t, batch_rule_t batch_rule>
  24601. at::Tensor row_indices_copy_generated_plumbing(const at::Tensor & self) {
  24602. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24603. auto maybe_layer = maybeCurrentDynamicLayer();
  24604. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24605. int64_t cur_level = maybe_layer->layerId();
  24606. if (!isBatchedAtLevel(self, cur_level)) {
  24607. return at::_ops::row_indices_copy::call(self);
  24608. }
  24609. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24610. auto results = batch_rule(self_value, self_bdim);
  24611. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24612. }
  24613. template <typename batch_rule_t, batch_rule_t batch_rule>
  24614. ::std::vector<at::Tensor> unbind_copy_int_generated_plumbing(const at::Tensor & self, int64_t dim) {
  24615. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24616. auto maybe_layer = maybeCurrentDynamicLayer();
  24617. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24618. int64_t cur_level = maybe_layer->layerId();
  24619. if (!isBatchedAtLevel(self, cur_level)) {
  24620. return at::_ops::unbind_copy_int::call(self, dim);
  24621. }
  24622. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24623. auto results = batch_rule(self_value, self_bdim, dim);
  24624. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  24625. }
  24626. template <typename batch_rule_t, batch_rule_t batch_rule>
  24627. void unbind_copy_int_out_generated_plumbing(const at::Tensor & self, int64_t dim, at::TensorList out) {
  24628. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24629. auto maybe_layer = maybeCurrentDynamicLayer();
  24630. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  24631. int64_t cur_level = maybe_layer->layerId();
  24632. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  24633. return at::_ops::unbind_copy_int_out::call(self, dim, out);
  24634. }
  24635. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24636. batch_rule(self_value, self_bdim, dim, out);
  24637. }
  24638. template <typename batch_rule_t, batch_rule_t batch_rule>
  24639. void split_copy_Tensor_out_generated_plumbing(const at::Tensor & self, c10::SymInt split_size, int64_t dim, at::TensorList out) {
  24640. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24641. auto maybe_layer = maybeCurrentDynamicLayer();
  24642. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  24643. int64_t cur_level = maybe_layer->layerId();
  24644. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  24645. return at::_ops::split_copy_Tensor_out::call(self, split_size, dim, out);
  24646. }
  24647. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24648. batch_rule(self_value, self_bdim, split_size, dim, out);
  24649. }
  24650. template <typename batch_rule_t, batch_rule_t batch_rule>
  24651. void split_with_sizes_copy_out_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef split_sizes, int64_t dim, at::TensorList out) {
  24652. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24653. auto maybe_layer = maybeCurrentDynamicLayer();
  24654. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  24655. int64_t cur_level = maybe_layer->layerId();
  24656. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  24657. return at::_ops::split_with_sizes_copy_out::call(self, split_sizes, dim, out);
  24658. }
  24659. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24660. batch_rule(self_value, self_bdim, split_sizes, dim, out);
  24661. }
  24662. template <typename batch_rule_t, batch_rule_t batch_rule>
  24663. at::Tensor view_copy_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef size) {
  24664. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24665. auto maybe_layer = maybeCurrentDynamicLayer();
  24666. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24667. int64_t cur_level = maybe_layer->layerId();
  24668. if (!isBatchedAtLevel(self, cur_level)) {
  24669. return at::_ops::view_copy::call(self, size);
  24670. }
  24671. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24672. auto results = batch_rule(self_value, self_bdim, size);
  24673. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24674. }
  24675. template <typename batch_rule_t, batch_rule_t batch_rule>
  24676. at::Tensor view_copy_dtype_generated_plumbing(const at::Tensor & self, at::ScalarType dtype) {
  24677. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24678. auto maybe_layer = maybeCurrentDynamicLayer();
  24679. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24680. int64_t cur_level = maybe_layer->layerId();
  24681. if (!isBatchedAtLevel(self, cur_level)) {
  24682. return at::_ops::view_copy_dtype::call(self, dtype);
  24683. }
  24684. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24685. auto results = batch_rule(self_value, self_bdim, dtype);
  24686. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24687. }
  24688. template <typename batch_rule_t, batch_rule_t batch_rule>
  24689. at::Tensor unfold_copy_generated_plumbing(const at::Tensor & self, int64_t dimension, int64_t size, int64_t step) {
  24690. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24691. auto maybe_layer = maybeCurrentDynamicLayer();
  24692. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24693. int64_t cur_level = maybe_layer->layerId();
  24694. if (!isBatchedAtLevel(self, cur_level)) {
  24695. return at::_ops::unfold_copy::call(self, dimension, size, step);
  24696. }
  24697. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24698. auto results = batch_rule(self_value, self_bdim, dimension, size, step);
  24699. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24700. }
  24701. template <typename batch_rule_t, batch_rule_t batch_rule>
  24702. at::Tensor alias_copy_generated_plumbing(const at::Tensor & self) {
  24703. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24704. auto maybe_layer = maybeCurrentDynamicLayer();
  24705. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24706. int64_t cur_level = maybe_layer->layerId();
  24707. if (!isBatchedAtLevel(self, cur_level)) {
  24708. return at::_ops::alias_copy::call(self);
  24709. }
  24710. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24711. auto results = batch_rule(self_value, self_bdim);
  24712. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24713. }
  24714. template <typename batch_rule_t, batch_rule_t batch_rule>
  24715. at::Tensor to_padded_tensor_generated_plumbing(const at::Tensor & self, double padding, at::OptionalSymIntArrayRef output_size) {
  24716. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24717. auto maybe_layer = maybeCurrentDynamicLayer();
  24718. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24719. int64_t cur_level = maybe_layer->layerId();
  24720. if (!isBatchedAtLevel(self, cur_level)) {
  24721. return at::_ops::to_padded_tensor::call(self, padding, output_size);
  24722. }
  24723. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24724. auto results = batch_rule(self_value, self_bdim, padding, output_size);
  24725. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24726. }
  24727. template <typename batch_rule_t, batch_rule_t batch_rule>
  24728. at::Tensor _jagged_to_padded_dense_forward_generated_plumbing(const at::Tensor & values, at::TensorList offsets, c10::SymIntArrayRef max_lengths, double padding_value) {
  24729. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24730. auto maybe_layer = maybeCurrentDynamicLayer();
  24731. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24732. int64_t cur_level = maybe_layer->layerId();
  24733. if (!isBatchedAtLevel(values, cur_level) && !isBatchedAtLevel(offsets, cur_level)) {
  24734. return at::_ops::_jagged_to_padded_dense_forward::call(values, offsets, max_lengths, padding_value);
  24735. }
  24736. auto [values_value, values_bdim] = unwrapTensorAtLevel(values, cur_level);
  24737. auto results = batch_rule(values_value, values_bdim, offsets, max_lengths, padding_value);
  24738. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24739. }
  24740. template <typename batch_rule_t, batch_rule_t batch_rule>
  24741. at::Tensor _padded_dense_to_jagged_forward_generated_plumbing(const at::Tensor & dense, at::TensorList offsets, ::std::optional<c10::SymInt> total_L) {
  24742. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24743. auto maybe_layer = maybeCurrentDynamicLayer();
  24744. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24745. int64_t cur_level = maybe_layer->layerId();
  24746. if (!isBatchedAtLevel(dense, cur_level) && !isBatchedAtLevel(offsets, cur_level)) {
  24747. return at::_ops::_padded_dense_to_jagged_forward::call(dense, offsets, total_L);
  24748. }
  24749. auto [dense_value, dense_bdim] = unwrapTensorAtLevel(dense, cur_level);
  24750. auto results = batch_rule(dense_value, dense_bdim, offsets, total_L);
  24751. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24752. }
  24753. template <typename batch_rule_t, batch_rule_t batch_rule>
  24754. at::Tensor _nested_from_padded_tensor_generated_plumbing(const at::Tensor & padded, const at::Tensor & offsets, const at::Tensor & dummy, int64_t ragged_idx, const ::std::optional<at::Tensor> & min_seqlen, const ::std::optional<at::Tensor> & max_seqlen, ::std::optional<c10::SymInt> sum_S) {
  24755. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24756. auto maybe_layer = maybeCurrentDynamicLayer();
  24757. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24758. int64_t cur_level = maybe_layer->layerId();
  24759. if (!isBatchedAtLevel(padded, cur_level) && !isBatchedAtLevel(offsets, cur_level) && !isBatchedAtLevel(dummy, cur_level) && !isBatchedAtLevel(min_seqlen, cur_level) && !isBatchedAtLevel(max_seqlen, cur_level)) {
  24760. return at::_ops::_nested_from_padded_tensor::call(padded, offsets, dummy, ragged_idx, min_seqlen, max_seqlen, sum_S);
  24761. }
  24762. auto [padded_value, padded_bdim] = unwrapTensorAtLevel(padded, cur_level);
  24763. auto [offsets_value, offsets_bdim] = unwrapTensorAtLevel(offsets, cur_level);
  24764. auto [dummy_value, dummy_bdim] = unwrapTensorAtLevel(dummy, cur_level);
  24765. std::optional<Tensor> min_seqlen_value;
  24766. std::optional<int64_t> min_seqlen_bdim;
  24767. if (min_seqlen) {
  24768. std::tie(min_seqlen_value, min_seqlen_bdim) = unwrapTensorAtLevel(min_seqlen.value(), cur_level);
  24769. }
  24770. std::optional<Tensor> max_seqlen_value;
  24771. std::optional<int64_t> max_seqlen_bdim;
  24772. if (max_seqlen) {
  24773. std::tie(max_seqlen_value, max_seqlen_bdim) = unwrapTensorAtLevel(max_seqlen.value(), cur_level);
  24774. }
  24775. auto results = batch_rule(padded_value, padded_bdim, offsets_value, offsets_bdim, dummy_value, dummy_bdim, ragged_idx, min_seqlen_value, min_seqlen_bdim, max_seqlen_value, max_seqlen_bdim, sum_S);
  24776. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24777. }
  24778. template <typename batch_rule_t, batch_rule_t batch_rule>
  24779. at::Tensor _nested_tensor_softmax_with_shape_generated_plumbing(const at::Tensor & self, const at::Tensor & query) {
  24780. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24781. auto maybe_layer = maybeCurrentDynamicLayer();
  24782. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24783. int64_t cur_level = maybe_layer->layerId();
  24784. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(query, cur_level)) {
  24785. return at::_ops::_nested_tensor_softmax_with_shape::call(self, query);
  24786. }
  24787. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24788. auto [query_value, query_bdim] = unwrapTensorAtLevel(query, cur_level);
  24789. auto results = batch_rule(self_value, self_bdim, query_value, query_bdim);
  24790. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24791. }
  24792. template <typename batch_rule_t, batch_rule_t batch_rule>
  24793. at::Tensor _safe_softmax_generated_plumbing(const at::Tensor & self, int64_t dim, ::std::optional<at::ScalarType> dtype) {
  24794. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24795. auto maybe_layer = maybeCurrentDynamicLayer();
  24796. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24797. int64_t cur_level = maybe_layer->layerId();
  24798. if (!isBatchedAtLevel(self, cur_level)) {
  24799. return at::_ops::_safe_softmax::call(self, dim, dtype);
  24800. }
  24801. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  24802. auto results = batch_rule(self_value, self_bdim, dim, dtype);
  24803. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24804. }
  24805. template <typename batch_rule_t, batch_rule_t batch_rule>
  24806. at::Tensor _transformer_encoder_layer_fwd_generated_plumbing(const at::Tensor & src, int64_t embed_dim, int64_t num_heads, const at::Tensor & qkv_weight, const at::Tensor & qkv_bias, const at::Tensor & proj_weight, const at::Tensor & proj_bias, bool use_gelu, bool norm_first, double eps, const at::Tensor & norm_weight_1, const at::Tensor & norm_bias_1, const at::Tensor & norm_weight_2, const at::Tensor & norm_bias_2, const at::Tensor & ffn_weight_1, const at::Tensor & ffn_bias_1, const at::Tensor & ffn_weight_2, const at::Tensor & ffn_bias_2, const ::std::optional<at::Tensor> & mask, ::std::optional<int64_t> mask_type) {
  24807. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24808. auto maybe_layer = maybeCurrentDynamicLayer();
  24809. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24810. int64_t cur_level = maybe_layer->layerId();
  24811. if (!isBatchedAtLevel(src, cur_level) && !isBatchedAtLevel(qkv_weight, cur_level) && !isBatchedAtLevel(qkv_bias, cur_level) && !isBatchedAtLevel(proj_weight, cur_level) && !isBatchedAtLevel(proj_bias, cur_level) && !isBatchedAtLevel(norm_weight_1, cur_level) && !isBatchedAtLevel(norm_bias_1, cur_level) && !isBatchedAtLevel(norm_weight_2, cur_level) && !isBatchedAtLevel(norm_bias_2, cur_level) && !isBatchedAtLevel(ffn_weight_1, cur_level) && !isBatchedAtLevel(ffn_bias_1, cur_level) && !isBatchedAtLevel(ffn_weight_2, cur_level) && !isBatchedAtLevel(ffn_bias_2, cur_level) && !isBatchedAtLevel(mask, cur_level)) {
  24812. return at::_ops::_transformer_encoder_layer_fwd::call(src, embed_dim, num_heads, qkv_weight, qkv_bias, proj_weight, proj_bias, use_gelu, norm_first, eps, norm_weight_1, norm_bias_1, norm_weight_2, norm_bias_2, ffn_weight_1, ffn_bias_1, ffn_weight_2, ffn_bias_2, mask, mask_type);
  24813. }
  24814. auto [src_value, src_bdim] = unwrapTensorAtLevel(src, cur_level);
  24815. auto [qkv_weight_value, qkv_weight_bdim] = unwrapTensorAtLevel(qkv_weight, cur_level);
  24816. auto [qkv_bias_value, qkv_bias_bdim] = unwrapTensorAtLevel(qkv_bias, cur_level);
  24817. auto [proj_weight_value, proj_weight_bdim] = unwrapTensorAtLevel(proj_weight, cur_level);
  24818. auto [proj_bias_value, proj_bias_bdim] = unwrapTensorAtLevel(proj_bias, cur_level);
  24819. auto [norm_weight_1_value, norm_weight_1_bdim] = unwrapTensorAtLevel(norm_weight_1, cur_level);
  24820. auto [norm_bias_1_value, norm_bias_1_bdim] = unwrapTensorAtLevel(norm_bias_1, cur_level);
  24821. auto [norm_weight_2_value, norm_weight_2_bdim] = unwrapTensorAtLevel(norm_weight_2, cur_level);
  24822. auto [norm_bias_2_value, norm_bias_2_bdim] = unwrapTensorAtLevel(norm_bias_2, cur_level);
  24823. auto [ffn_weight_1_value, ffn_weight_1_bdim] = unwrapTensorAtLevel(ffn_weight_1, cur_level);
  24824. auto [ffn_bias_1_value, ffn_bias_1_bdim] = unwrapTensorAtLevel(ffn_bias_1, cur_level);
  24825. auto [ffn_weight_2_value, ffn_weight_2_bdim] = unwrapTensorAtLevel(ffn_weight_2, cur_level);
  24826. auto [ffn_bias_2_value, ffn_bias_2_bdim] = unwrapTensorAtLevel(ffn_bias_2, cur_level);
  24827. std::optional<Tensor> mask_value;
  24828. std::optional<int64_t> mask_bdim;
  24829. if (mask) {
  24830. std::tie(mask_value, mask_bdim) = unwrapTensorAtLevel(mask.value(), cur_level);
  24831. }
  24832. auto results = batch_rule(src_value, src_bdim, embed_dim, num_heads, qkv_weight_value, qkv_weight_bdim, qkv_bias_value, qkv_bias_bdim, proj_weight_value, proj_weight_bdim, proj_bias_value, proj_bias_bdim, use_gelu, norm_first, eps, norm_weight_1_value, norm_weight_1_bdim, norm_bias_1_value, norm_bias_1_bdim, norm_weight_2_value, norm_weight_2_bdim, norm_bias_2_value, norm_bias_2_bdim, ffn_weight_1_value, ffn_weight_1_bdim, ffn_bias_1_value, ffn_bias_1_bdim, ffn_weight_2_value, ffn_weight_2_bdim, ffn_bias_2_value, ffn_bias_2_bdim, mask_value, mask_bdim, mask_type);
  24833. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24834. }
  24835. template <typename batch_rule_t, batch_rule_t batch_rule>
  24836. ::std::tuple<at::Tensor,at::Tensor> _native_multi_head_attention_generated_plumbing(const at::Tensor & query, const at::Tensor & key, const at::Tensor & value, int64_t embed_dim, int64_t num_head, const at::Tensor & qkv_weight, const at::Tensor & qkv_bias, const at::Tensor & proj_weight, const at::Tensor & proj_bias, const ::std::optional<at::Tensor> & mask, bool need_weights, bool average_attn_weights, ::std::optional<int64_t> mask_type) {
  24837. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24838. auto maybe_layer = maybeCurrentDynamicLayer();
  24839. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24840. int64_t cur_level = maybe_layer->layerId();
  24841. if (!isBatchedAtLevel(query, cur_level) && !isBatchedAtLevel(key, cur_level) && !isBatchedAtLevel(value, cur_level) && !isBatchedAtLevel(qkv_weight, cur_level) && !isBatchedAtLevel(qkv_bias, cur_level) && !isBatchedAtLevel(proj_weight, cur_level) && !isBatchedAtLevel(proj_bias, cur_level) && !isBatchedAtLevel(mask, cur_level)) {
  24842. return at::_ops::_native_multi_head_attention::call(query, key, value, embed_dim, num_head, qkv_weight, qkv_bias, proj_weight, proj_bias, mask, need_weights, average_attn_weights, mask_type);
  24843. }
  24844. auto [query_value, query_bdim] = unwrapTensorAtLevel(query, cur_level);
  24845. auto [key_value, key_bdim] = unwrapTensorAtLevel(key, cur_level);
  24846. auto [value_value, value_bdim] = unwrapTensorAtLevel(value, cur_level);
  24847. auto [qkv_weight_value, qkv_weight_bdim] = unwrapTensorAtLevel(qkv_weight, cur_level);
  24848. auto [qkv_bias_value, qkv_bias_bdim] = unwrapTensorAtLevel(qkv_bias, cur_level);
  24849. auto [proj_weight_value, proj_weight_bdim] = unwrapTensorAtLevel(proj_weight, cur_level);
  24850. auto [proj_bias_value, proj_bias_bdim] = unwrapTensorAtLevel(proj_bias, cur_level);
  24851. std::optional<Tensor> mask_value;
  24852. std::optional<int64_t> mask_bdim;
  24853. if (mask) {
  24854. std::tie(mask_value, mask_bdim) = unwrapTensorAtLevel(mask.value(), cur_level);
  24855. }
  24856. auto results = batch_rule(query_value, query_bdim, key_value, key_bdim, value_value, value_bdim, embed_dim, num_head, qkv_weight_value, qkv_weight_bdim, qkv_bias_value, qkv_bias_bdim, proj_weight_value, proj_weight_bdim, proj_bias_value, proj_bias_bdim, mask_value, mask_bdim, need_weights, average_attn_weights, mask_type);
  24857. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  24858. }
  24859. template <typename batch_rule_t, batch_rule_t batch_rule>
  24860. at::Tensor scaled_dot_product_attention_generated_plumbing(const at::Tensor & query, const at::Tensor & key, const at::Tensor & value, const ::std::optional<at::Tensor> & attn_mask, double dropout_p, bool is_causal, ::std::optional<double> scale, bool enable_gqa) {
  24861. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24862. auto maybe_layer = maybeCurrentDynamicLayer();
  24863. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24864. int64_t cur_level = maybe_layer->layerId();
  24865. if (!isBatchedAtLevel(query, cur_level) && !isBatchedAtLevel(key, cur_level) && !isBatchedAtLevel(value, cur_level) && !isBatchedAtLevel(attn_mask, cur_level)) {
  24866. return at::_ops::scaled_dot_product_attention::call(query, key, value, attn_mask, dropout_p, is_causal, scale, enable_gqa);
  24867. }
  24868. auto [query_value, query_bdim] = unwrapTensorAtLevel(query, cur_level);
  24869. auto [key_value, key_bdim] = unwrapTensorAtLevel(key, cur_level);
  24870. auto [value_value, value_bdim] = unwrapTensorAtLevel(value, cur_level);
  24871. std::optional<Tensor> attn_mask_value;
  24872. std::optional<int64_t> attn_mask_bdim;
  24873. if (attn_mask) {
  24874. std::tie(attn_mask_value, attn_mask_bdim) = unwrapTensorAtLevel(attn_mask.value(), cur_level);
  24875. }
  24876. auto results = batch_rule(query_value, query_bdim, key_value, key_bdim, value_value, value_bdim, attn_mask_value, attn_mask_bdim, dropout_p, is_causal, scale, enable_gqa);
  24877. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  24878. }
  24879. template <typename batch_rule_t, batch_rule_t batch_rule>
  24880. ::std::tuple<at::Tensor,at::Tensor> _scaled_dot_product_attention_math_generated_plumbing(const at::Tensor & query, const at::Tensor & key, const at::Tensor & value, const ::std::optional<at::Tensor> & attn_mask, double dropout_p, bool is_causal, const ::std::optional<at::Tensor> & dropout_mask, ::std::optional<double> scale, bool enable_gqa) {
  24881. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24882. auto maybe_layer = maybeCurrentDynamicLayer();
  24883. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24884. int64_t cur_level = maybe_layer->layerId();
  24885. if (!isBatchedAtLevel(query, cur_level) && !isBatchedAtLevel(key, cur_level) && !isBatchedAtLevel(value, cur_level) && !isBatchedAtLevel(attn_mask, cur_level) && !isBatchedAtLevel(dropout_mask, cur_level)) {
  24886. return at::_ops::_scaled_dot_product_attention_math::call(query, key, value, attn_mask, dropout_p, is_causal, dropout_mask, scale, enable_gqa);
  24887. }
  24888. auto [query_value, query_bdim] = unwrapTensorAtLevel(query, cur_level);
  24889. auto [key_value, key_bdim] = unwrapTensorAtLevel(key, cur_level);
  24890. auto [value_value, value_bdim] = unwrapTensorAtLevel(value, cur_level);
  24891. std::optional<Tensor> attn_mask_value;
  24892. std::optional<int64_t> attn_mask_bdim;
  24893. if (attn_mask) {
  24894. std::tie(attn_mask_value, attn_mask_bdim) = unwrapTensorAtLevel(attn_mask.value(), cur_level);
  24895. }
  24896. std::optional<Tensor> dropout_mask_value;
  24897. std::optional<int64_t> dropout_mask_bdim;
  24898. if (dropout_mask) {
  24899. std::tie(dropout_mask_value, dropout_mask_bdim) = unwrapTensorAtLevel(dropout_mask.value(), cur_level);
  24900. }
  24901. auto results = batch_rule(query_value, query_bdim, key_value, key_bdim, value_value, value_bdim, attn_mask_value, attn_mask_bdim, dropout_p, is_causal, dropout_mask_value, dropout_mask_bdim, scale, enable_gqa);
  24902. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  24903. }
  24904. template <typename batch_rule_t, batch_rule_t batch_rule>
  24905. ::std::tuple<at::Tensor,at::Tensor> _scaled_dot_product_attention_math_for_mps_generated_plumbing(const at::Tensor & query, const at::Tensor & key, const at::Tensor & value, const ::std::optional<at::Tensor> & attn_mask, double dropout_p, bool is_causal, const ::std::optional<at::Tensor> & dropout_mask, ::std::optional<double> scale) {
  24906. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24907. auto maybe_layer = maybeCurrentDynamicLayer();
  24908. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24909. int64_t cur_level = maybe_layer->layerId();
  24910. if (!isBatchedAtLevel(query, cur_level) && !isBatchedAtLevel(key, cur_level) && !isBatchedAtLevel(value, cur_level) && !isBatchedAtLevel(attn_mask, cur_level) && !isBatchedAtLevel(dropout_mask, cur_level)) {
  24911. return at::_ops::_scaled_dot_product_attention_math_for_mps::call(query, key, value, attn_mask, dropout_p, is_causal, dropout_mask, scale);
  24912. }
  24913. auto [query_value, query_bdim] = unwrapTensorAtLevel(query, cur_level);
  24914. auto [key_value, key_bdim] = unwrapTensorAtLevel(key, cur_level);
  24915. auto [value_value, value_bdim] = unwrapTensorAtLevel(value, cur_level);
  24916. std::optional<Tensor> attn_mask_value;
  24917. std::optional<int64_t> attn_mask_bdim;
  24918. if (attn_mask) {
  24919. std::tie(attn_mask_value, attn_mask_bdim) = unwrapTensorAtLevel(attn_mask.value(), cur_level);
  24920. }
  24921. std::optional<Tensor> dropout_mask_value;
  24922. std::optional<int64_t> dropout_mask_bdim;
  24923. if (dropout_mask) {
  24924. std::tie(dropout_mask_value, dropout_mask_bdim) = unwrapTensorAtLevel(dropout_mask.value(), cur_level);
  24925. }
  24926. auto results = batch_rule(query_value, query_bdim, key_value, key_bdim, value_value, value_bdim, attn_mask_value, attn_mask_bdim, dropout_p, is_causal, dropout_mask_value, dropout_mask_bdim, scale);
  24927. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  24928. }
  24929. template <typename batch_rule_t, batch_rule_t batch_rule>
  24930. ::std::tuple<at::Tensor,at::Tensor,at::Tensor,at::Tensor,c10::SymInt,c10::SymInt,at::Tensor,at::Tensor,at::Tensor> _scaled_dot_product_flash_attention_generated_plumbing(const at::Tensor & query, const at::Tensor & key, const at::Tensor & value, double dropout_p, bool is_causal, bool return_debug_mask, ::std::optional<double> scale) {
  24931. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24932. auto maybe_layer = maybeCurrentDynamicLayer();
  24933. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24934. int64_t cur_level = maybe_layer->layerId();
  24935. if (!isBatchedAtLevel(query, cur_level) && !isBatchedAtLevel(key, cur_level) && !isBatchedAtLevel(value, cur_level)) {
  24936. return at::_ops::_scaled_dot_product_flash_attention::call(query, key, value, dropout_p, is_causal, return_debug_mask, scale);
  24937. }
  24938. auto [query_value, query_bdim] = unwrapTensorAtLevel(query, cur_level);
  24939. auto [key_value, key_bdim] = unwrapTensorAtLevel(key, cur_level);
  24940. auto [value_value, value_bdim] = unwrapTensorAtLevel(value, cur_level);
  24941. auto results = batch_rule(query_value, query_bdim, key_value, key_bdim, value_value, value_bdim, dropout_p, is_causal, return_debug_mask, scale);
  24942. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level), makeBatched(std::get<6>(results), std::get<7>(results), cur_level), std::get<8>(results), std::get<9>(results), makeBatched(std::get<10>(results), std::get<11>(results), cur_level), makeBatched(std::get<12>(results), std::get<13>(results), cur_level), makeBatched(std::get<14>(results), std::get<15>(results), cur_level));
  24943. }
  24944. template <typename batch_rule_t, batch_rule_t batch_rule>
  24945. ::std::tuple<at::Tensor,at::Tensor> _scaled_dot_product_flash_attention_for_cpu_generated_plumbing(const at::Tensor & query, const at::Tensor & key, const at::Tensor & value, double dropout_p, bool is_causal, const ::std::optional<at::Tensor> & attn_mask, ::std::optional<double> scale) {
  24946. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24947. auto maybe_layer = maybeCurrentDynamicLayer();
  24948. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24949. int64_t cur_level = maybe_layer->layerId();
  24950. if (!isBatchedAtLevel(query, cur_level) && !isBatchedAtLevel(key, cur_level) && !isBatchedAtLevel(value, cur_level) && !isBatchedAtLevel(attn_mask, cur_level)) {
  24951. return at::_ops::_scaled_dot_product_flash_attention_for_cpu::call(query, key, value, dropout_p, is_causal, attn_mask, scale);
  24952. }
  24953. auto [query_value, query_bdim] = unwrapTensorAtLevel(query, cur_level);
  24954. auto [key_value, key_bdim] = unwrapTensorAtLevel(key, cur_level);
  24955. auto [value_value, value_bdim] = unwrapTensorAtLevel(value, cur_level);
  24956. std::optional<Tensor> attn_mask_value;
  24957. std::optional<int64_t> attn_mask_bdim;
  24958. if (attn_mask) {
  24959. std::tie(attn_mask_value, attn_mask_bdim) = unwrapTensorAtLevel(attn_mask.value(), cur_level);
  24960. }
  24961. auto results = batch_rule(query_value, query_bdim, key_value, key_bdim, value_value, value_bdim, dropout_p, is_causal, attn_mask_value, attn_mask_bdim, scale);
  24962. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  24963. }
  24964. template <typename batch_rule_t, batch_rule_t batch_rule>
  24965. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> _scaled_dot_product_flash_attention_backward_generated_plumbing(const at::Tensor & grad_out, const at::Tensor & query, const at::Tensor & key, const at::Tensor & value, const at::Tensor & out, const at::Tensor & logsumexp, const at::Tensor & cum_seq_q, const at::Tensor & cum_seq_k, c10::SymInt max_q, c10::SymInt max_k, double dropout_p, bool is_causal, const at::Tensor & philox_seed, const at::Tensor & philox_offset, ::std::optional<double> scale) {
  24966. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24967. auto maybe_layer = maybeCurrentDynamicLayer();
  24968. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24969. int64_t cur_level = maybe_layer->layerId();
  24970. if (!isBatchedAtLevel(grad_out, cur_level) && !isBatchedAtLevel(query, cur_level) && !isBatchedAtLevel(key, cur_level) && !isBatchedAtLevel(value, cur_level) && !isBatchedAtLevel(out, cur_level) && !isBatchedAtLevel(logsumexp, cur_level) && !isBatchedAtLevel(cum_seq_q, cur_level) && !isBatchedAtLevel(cum_seq_k, cur_level) && !isBatchedAtLevel(philox_seed, cur_level) && !isBatchedAtLevel(philox_offset, cur_level)) {
  24971. return at::_ops::_scaled_dot_product_flash_attention_backward::call(grad_out, query, key, value, out, logsumexp, cum_seq_q, cum_seq_k, max_q, max_k, dropout_p, is_causal, philox_seed, philox_offset, scale);
  24972. }
  24973. auto [grad_out_value, grad_out_bdim] = unwrapTensorAtLevel(grad_out, cur_level);
  24974. auto [query_value, query_bdim] = unwrapTensorAtLevel(query, cur_level);
  24975. auto [key_value, key_bdim] = unwrapTensorAtLevel(key, cur_level);
  24976. auto [value_value, value_bdim] = unwrapTensorAtLevel(value, cur_level);
  24977. auto [out_value, out_bdim] = unwrapTensorAtLevel(out, cur_level);
  24978. auto [logsumexp_value, logsumexp_bdim] = unwrapTensorAtLevel(logsumexp, cur_level);
  24979. auto [cum_seq_q_value, cum_seq_q_bdim] = unwrapTensorAtLevel(cum_seq_q, cur_level);
  24980. auto [cum_seq_k_value, cum_seq_k_bdim] = unwrapTensorAtLevel(cum_seq_k, cur_level);
  24981. auto [philox_seed_value, philox_seed_bdim] = unwrapTensorAtLevel(philox_seed, cur_level);
  24982. auto [philox_offset_value, philox_offset_bdim] = unwrapTensorAtLevel(philox_offset, cur_level);
  24983. auto results = batch_rule(grad_out_value, grad_out_bdim, query_value, query_bdim, key_value, key_bdim, value_value, value_bdim, out_value, out_bdim, logsumexp_value, logsumexp_bdim, cum_seq_q_value, cum_seq_q_bdim, cum_seq_k_value, cum_seq_k_bdim, max_q, max_k, dropout_p, is_causal, philox_seed_value, philox_seed_bdim, philox_offset_value, philox_offset_bdim, scale);
  24984. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  24985. }
  24986. template <typename batch_rule_t, batch_rule_t batch_rule>
  24987. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> _scaled_dot_product_flash_attention_for_cpu_backward_generated_plumbing(const at::Tensor & grad_out, const at::Tensor & query, const at::Tensor & key, const at::Tensor & value, const at::Tensor & out, const at::Tensor & logsumexp, double dropout_p, bool is_causal, const ::std::optional<at::Tensor> & attn_mask, ::std::optional<double> scale) {
  24988. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  24989. auto maybe_layer = maybeCurrentDynamicLayer();
  24990. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  24991. int64_t cur_level = maybe_layer->layerId();
  24992. if (!isBatchedAtLevel(grad_out, cur_level) && !isBatchedAtLevel(query, cur_level) && !isBatchedAtLevel(key, cur_level) && !isBatchedAtLevel(value, cur_level) && !isBatchedAtLevel(out, cur_level) && !isBatchedAtLevel(logsumexp, cur_level) && !isBatchedAtLevel(attn_mask, cur_level)) {
  24993. return at::_ops::_scaled_dot_product_flash_attention_for_cpu_backward::call(grad_out, query, key, value, out, logsumexp, dropout_p, is_causal, attn_mask, scale);
  24994. }
  24995. auto [grad_out_value, grad_out_bdim] = unwrapTensorAtLevel(grad_out, cur_level);
  24996. auto [query_value, query_bdim] = unwrapTensorAtLevel(query, cur_level);
  24997. auto [key_value, key_bdim] = unwrapTensorAtLevel(key, cur_level);
  24998. auto [value_value, value_bdim] = unwrapTensorAtLevel(value, cur_level);
  24999. auto [out_value, out_bdim] = unwrapTensorAtLevel(out, cur_level);
  25000. auto [logsumexp_value, logsumexp_bdim] = unwrapTensorAtLevel(logsumexp, cur_level);
  25001. std::optional<Tensor> attn_mask_value;
  25002. std::optional<int64_t> attn_mask_bdim;
  25003. if (attn_mask) {
  25004. std::tie(attn_mask_value, attn_mask_bdim) = unwrapTensorAtLevel(attn_mask.value(), cur_level);
  25005. }
  25006. auto results = batch_rule(grad_out_value, grad_out_bdim, query_value, query_bdim, key_value, key_bdim, value_value, value_bdim, out_value, out_bdim, logsumexp_value, logsumexp_bdim, dropout_p, is_causal, attn_mask_value, attn_mask_bdim, scale);
  25007. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  25008. }
  25009. template <typename batch_rule_t, batch_rule_t batch_rule>
  25010. ::std::tuple<at::Tensor,at::Tensor,at::Tensor,at::Tensor> _scaled_dot_product_fused_attention_overrideable_backward_generated_plumbing(const at::Tensor & grad_out, const at::Tensor & query, const at::Tensor & key, const at::Tensor & value, const at::Tensor & attn_bias, ::std::array<bool,4> grad_input_mask, const at::Tensor & out, const at::Tensor & logsumexp, const at::Tensor & cum_seq_q, const at::Tensor & cum_seq_k, c10::SymInt max_q, c10::SymInt max_k, double dropout_p, bool is_causal, const at::Tensor & philox_seed, const at::Tensor & philox_offset, ::std::optional<double> scale) {
  25011. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25012. auto maybe_layer = maybeCurrentDynamicLayer();
  25013. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25014. int64_t cur_level = maybe_layer->layerId();
  25015. if (!isBatchedAtLevel(grad_out, cur_level) && !isBatchedAtLevel(query, cur_level) && !isBatchedAtLevel(key, cur_level) && !isBatchedAtLevel(value, cur_level) && !isBatchedAtLevel(attn_bias, cur_level) && !isBatchedAtLevel(out, cur_level) && !isBatchedAtLevel(logsumexp, cur_level) && !isBatchedAtLevel(cum_seq_q, cur_level) && !isBatchedAtLevel(cum_seq_k, cur_level) && !isBatchedAtLevel(philox_seed, cur_level) && !isBatchedAtLevel(philox_offset, cur_level)) {
  25016. return at::_ops::_scaled_dot_product_fused_attention_overrideable_backward::call(grad_out, query, key, value, attn_bias, grad_input_mask, out, logsumexp, cum_seq_q, cum_seq_k, max_q, max_k, dropout_p, is_causal, philox_seed, philox_offset, scale);
  25017. }
  25018. auto [grad_out_value, grad_out_bdim] = unwrapTensorAtLevel(grad_out, cur_level);
  25019. auto [query_value, query_bdim] = unwrapTensorAtLevel(query, cur_level);
  25020. auto [key_value, key_bdim] = unwrapTensorAtLevel(key, cur_level);
  25021. auto [value_value, value_bdim] = unwrapTensorAtLevel(value, cur_level);
  25022. auto [attn_bias_value, attn_bias_bdim] = unwrapTensorAtLevel(attn_bias, cur_level);
  25023. auto [out_value, out_bdim] = unwrapTensorAtLevel(out, cur_level);
  25024. auto [logsumexp_value, logsumexp_bdim] = unwrapTensorAtLevel(logsumexp, cur_level);
  25025. auto [cum_seq_q_value, cum_seq_q_bdim] = unwrapTensorAtLevel(cum_seq_q, cur_level);
  25026. auto [cum_seq_k_value, cum_seq_k_bdim] = unwrapTensorAtLevel(cum_seq_k, cur_level);
  25027. auto [philox_seed_value, philox_seed_bdim] = unwrapTensorAtLevel(philox_seed, cur_level);
  25028. auto [philox_offset_value, philox_offset_bdim] = unwrapTensorAtLevel(philox_offset, cur_level);
  25029. auto results = batch_rule(grad_out_value, grad_out_bdim, query_value, query_bdim, key_value, key_bdim, value_value, value_bdim, attn_bias_value, attn_bias_bdim, grad_input_mask, out_value, out_bdim, logsumexp_value, logsumexp_bdim, cum_seq_q_value, cum_seq_q_bdim, cum_seq_k_value, cum_seq_k_bdim, max_q, max_k, dropout_p, is_causal, philox_seed_value, philox_seed_bdim, philox_offset_value, philox_offset_bdim, scale);
  25030. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level), makeBatched(std::get<6>(results), std::get<7>(results), cur_level));
  25031. }
  25032. template <typename batch_rule_t, batch_rule_t batch_rule>
  25033. ::std::tuple<at::Tensor,at::Tensor,at::Tensor,at::Tensor> _scaled_dot_product_efficient_attention_generated_plumbing(const at::Tensor & query, const at::Tensor & key, const at::Tensor & value, const ::std::optional<at::Tensor> & attn_bias, bool compute_log_sumexp, double dropout_p, bool is_causal, ::std::optional<double> scale) {
  25034. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25035. auto maybe_layer = maybeCurrentDynamicLayer();
  25036. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25037. int64_t cur_level = maybe_layer->layerId();
  25038. if (!isBatchedAtLevel(query, cur_level) && !isBatchedAtLevel(key, cur_level) && !isBatchedAtLevel(value, cur_level) && !isBatchedAtLevel(attn_bias, cur_level)) {
  25039. return at::_ops::_scaled_dot_product_efficient_attention::call(query, key, value, attn_bias, compute_log_sumexp, dropout_p, is_causal, scale);
  25040. }
  25041. auto [query_value, query_bdim] = unwrapTensorAtLevel(query, cur_level);
  25042. auto [key_value, key_bdim] = unwrapTensorAtLevel(key, cur_level);
  25043. auto [value_value, value_bdim] = unwrapTensorAtLevel(value, cur_level);
  25044. std::optional<Tensor> attn_bias_value;
  25045. std::optional<int64_t> attn_bias_bdim;
  25046. if (attn_bias) {
  25047. std::tie(attn_bias_value, attn_bias_bdim) = unwrapTensorAtLevel(attn_bias.value(), cur_level);
  25048. }
  25049. auto results = batch_rule(query_value, query_bdim, key_value, key_bdim, value_value, value_bdim, attn_bias_value, attn_bias_bdim, compute_log_sumexp, dropout_p, is_causal, scale);
  25050. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level), makeBatched(std::get<6>(results), std::get<7>(results), cur_level));
  25051. }
  25052. template <typename batch_rule_t, batch_rule_t batch_rule>
  25053. ::std::tuple<at::Tensor,at::Tensor,at::Tensor,at::Tensor> _scaled_dot_product_efficient_attention_backward_generated_plumbing(const at::Tensor & grad_out_, const at::Tensor & query, const at::Tensor & key, const at::Tensor & value, const at::Tensor & attn_bias, const at::Tensor & out, const at::Tensor & logsumexp, const at::Tensor & philox_seed, const at::Tensor & philox_offset, double dropout_p, ::std::array<bool,4> grad_input_mask, bool is_causal, ::std::optional<double> scale) {
  25054. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25055. auto maybe_layer = maybeCurrentDynamicLayer();
  25056. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25057. int64_t cur_level = maybe_layer->layerId();
  25058. if (!isBatchedAtLevel(grad_out_, cur_level) && !isBatchedAtLevel(query, cur_level) && !isBatchedAtLevel(key, cur_level) && !isBatchedAtLevel(value, cur_level) && !isBatchedAtLevel(attn_bias, cur_level) && !isBatchedAtLevel(out, cur_level) && !isBatchedAtLevel(logsumexp, cur_level) && !isBatchedAtLevel(philox_seed, cur_level) && !isBatchedAtLevel(philox_offset, cur_level)) {
  25059. return at::_ops::_scaled_dot_product_efficient_attention_backward::call(grad_out_, query, key, value, attn_bias, out, logsumexp, philox_seed, philox_offset, dropout_p, grad_input_mask, is_causal, scale);
  25060. }
  25061. auto [grad_out__value, grad_out__bdim] = unwrapTensorAtLevel(grad_out_, cur_level);
  25062. auto [query_value, query_bdim] = unwrapTensorAtLevel(query, cur_level);
  25063. auto [key_value, key_bdim] = unwrapTensorAtLevel(key, cur_level);
  25064. auto [value_value, value_bdim] = unwrapTensorAtLevel(value, cur_level);
  25065. auto [attn_bias_value, attn_bias_bdim] = unwrapTensorAtLevel(attn_bias, cur_level);
  25066. auto [out_value, out_bdim] = unwrapTensorAtLevel(out, cur_level);
  25067. auto [logsumexp_value, logsumexp_bdim] = unwrapTensorAtLevel(logsumexp, cur_level);
  25068. auto [philox_seed_value, philox_seed_bdim] = unwrapTensorAtLevel(philox_seed, cur_level);
  25069. auto [philox_offset_value, philox_offset_bdim] = unwrapTensorAtLevel(philox_offset, cur_level);
  25070. auto results = batch_rule(grad_out__value, grad_out__bdim, query_value, query_bdim, key_value, key_bdim, value_value, value_bdim, attn_bias_value, attn_bias_bdim, out_value, out_bdim, logsumexp_value, logsumexp_bdim, philox_seed_value, philox_seed_bdim, philox_offset_value, philox_offset_bdim, dropout_p, grad_input_mask, is_causal, scale);
  25071. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level), makeBatched(std::get<6>(results), std::get<7>(results), cur_level));
  25072. }
  25073. template <typename batch_rule_t, batch_rule_t batch_rule>
  25074. ::std::tuple<at::Tensor,at::Tensor,at::Tensor,at::Tensor,c10::SymInt,c10::SymInt,at::Tensor,at::Tensor,at::Tensor> _scaled_dot_product_cudnn_attention_generated_plumbing(const at::Tensor & query, const at::Tensor & key, const at::Tensor & value, const ::std::optional<at::Tensor> & attn_bias, bool compute_log_sumexp, double dropout_p, bool is_causal, bool return_debug_mask, ::std::optional<double> scale) {
  25075. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25076. auto maybe_layer = maybeCurrentDynamicLayer();
  25077. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25078. int64_t cur_level = maybe_layer->layerId();
  25079. if (!isBatchedAtLevel(query, cur_level) && !isBatchedAtLevel(key, cur_level) && !isBatchedAtLevel(value, cur_level) && !isBatchedAtLevel(attn_bias, cur_level)) {
  25080. return at::_ops::_scaled_dot_product_cudnn_attention::call(query, key, value, attn_bias, compute_log_sumexp, dropout_p, is_causal, return_debug_mask, scale);
  25081. }
  25082. auto [query_value, query_bdim] = unwrapTensorAtLevel(query, cur_level);
  25083. auto [key_value, key_bdim] = unwrapTensorAtLevel(key, cur_level);
  25084. auto [value_value, value_bdim] = unwrapTensorAtLevel(value, cur_level);
  25085. std::optional<Tensor> attn_bias_value;
  25086. std::optional<int64_t> attn_bias_bdim;
  25087. if (attn_bias) {
  25088. std::tie(attn_bias_value, attn_bias_bdim) = unwrapTensorAtLevel(attn_bias.value(), cur_level);
  25089. }
  25090. auto results = batch_rule(query_value, query_bdim, key_value, key_bdim, value_value, value_bdim, attn_bias_value, attn_bias_bdim, compute_log_sumexp, dropout_p, is_causal, return_debug_mask, scale);
  25091. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level), makeBatched(std::get<6>(results), std::get<7>(results), cur_level), std::get<8>(results), std::get<9>(results), makeBatched(std::get<10>(results), std::get<11>(results), cur_level), makeBatched(std::get<12>(results), std::get<13>(results), cur_level), makeBatched(std::get<14>(results), std::get<15>(results), cur_level));
  25092. }
  25093. template <typename batch_rule_t, batch_rule_t batch_rule>
  25094. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> _scaled_dot_product_cudnn_attention_backward_generated_plumbing(const at::Tensor & grad_out, const at::Tensor & query, const at::Tensor & key, const at::Tensor & value, const at::Tensor & out, const at::Tensor & logsumexp, const at::Tensor & philox_seed, const at::Tensor & philox_offset, const at::Tensor & attn_bias, const at::Tensor & cum_seq_q, const at::Tensor & cum_seq_k, c10::SymInt max_q, c10::SymInt max_k, double dropout_p, bool is_causal, ::std::optional<double> scale) {
  25095. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25096. auto maybe_layer = maybeCurrentDynamicLayer();
  25097. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25098. int64_t cur_level = maybe_layer->layerId();
  25099. if (!isBatchedAtLevel(grad_out, cur_level) && !isBatchedAtLevel(query, cur_level) && !isBatchedAtLevel(key, cur_level) && !isBatchedAtLevel(value, cur_level) && !isBatchedAtLevel(out, cur_level) && !isBatchedAtLevel(logsumexp, cur_level) && !isBatchedAtLevel(philox_seed, cur_level) && !isBatchedAtLevel(philox_offset, cur_level) && !isBatchedAtLevel(attn_bias, cur_level) && !isBatchedAtLevel(cum_seq_q, cur_level) && !isBatchedAtLevel(cum_seq_k, cur_level)) {
  25100. return at::_ops::_scaled_dot_product_cudnn_attention_backward::call(grad_out, query, key, value, out, logsumexp, philox_seed, philox_offset, attn_bias, cum_seq_q, cum_seq_k, max_q, max_k, dropout_p, is_causal, scale);
  25101. }
  25102. auto [grad_out_value, grad_out_bdim] = unwrapTensorAtLevel(grad_out, cur_level);
  25103. auto [query_value, query_bdim] = unwrapTensorAtLevel(query, cur_level);
  25104. auto [key_value, key_bdim] = unwrapTensorAtLevel(key, cur_level);
  25105. auto [value_value, value_bdim] = unwrapTensorAtLevel(value, cur_level);
  25106. auto [out_value, out_bdim] = unwrapTensorAtLevel(out, cur_level);
  25107. auto [logsumexp_value, logsumexp_bdim] = unwrapTensorAtLevel(logsumexp, cur_level);
  25108. auto [philox_seed_value, philox_seed_bdim] = unwrapTensorAtLevel(philox_seed, cur_level);
  25109. auto [philox_offset_value, philox_offset_bdim] = unwrapTensorAtLevel(philox_offset, cur_level);
  25110. auto [attn_bias_value, attn_bias_bdim] = unwrapTensorAtLevel(attn_bias, cur_level);
  25111. auto [cum_seq_q_value, cum_seq_q_bdim] = unwrapTensorAtLevel(cum_seq_q, cur_level);
  25112. auto [cum_seq_k_value, cum_seq_k_bdim] = unwrapTensorAtLevel(cum_seq_k, cur_level);
  25113. auto results = batch_rule(grad_out_value, grad_out_bdim, query_value, query_bdim, key_value, key_bdim, value_value, value_bdim, out_value, out_bdim, logsumexp_value, logsumexp_bdim, philox_seed_value, philox_seed_bdim, philox_offset_value, philox_offset_bdim, attn_bias_value, attn_bias_bdim, cum_seq_q_value, cum_seq_q_bdim, cum_seq_k_value, cum_seq_k_bdim, max_q, max_k, dropout_p, is_causal, scale);
  25114. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  25115. }
  25116. template <typename batch_rule_t, batch_rule_t batch_rule>
  25117. ::std::tuple<at::Tensor,at::Tensor,at::Tensor,at::Tensor,at::Tensor> _flash_attention_forward_generated_plumbing(const at::Tensor & query, const at::Tensor & key, const at::Tensor & value, const ::std::optional<at::Tensor> & cum_seq_q, const ::std::optional<at::Tensor> & cum_seq_k, c10::SymInt max_q, c10::SymInt max_k, double dropout_p, bool is_causal, bool return_debug_mask, ::std::optional<double> scale, ::std::optional<c10::SymInt> window_size_left, ::std::optional<c10::SymInt> window_size_right, const ::std::optional<at::Tensor> & seqused_k, const ::std::optional<at::Tensor> & alibi_slopes) {
  25118. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25119. auto maybe_layer = maybeCurrentDynamicLayer();
  25120. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25121. int64_t cur_level = maybe_layer->layerId();
  25122. if (!isBatchedAtLevel(query, cur_level) && !isBatchedAtLevel(key, cur_level) && !isBatchedAtLevel(value, cur_level) && !isBatchedAtLevel(cum_seq_q, cur_level) && !isBatchedAtLevel(cum_seq_k, cur_level) && !isBatchedAtLevel(seqused_k, cur_level) && !isBatchedAtLevel(alibi_slopes, cur_level)) {
  25123. return at::_ops::_flash_attention_forward::call(query, key, value, cum_seq_q, cum_seq_k, max_q, max_k, dropout_p, is_causal, return_debug_mask, scale, window_size_left, window_size_right, seqused_k, alibi_slopes);
  25124. }
  25125. auto [query_value, query_bdim] = unwrapTensorAtLevel(query, cur_level);
  25126. auto [key_value, key_bdim] = unwrapTensorAtLevel(key, cur_level);
  25127. auto [value_value, value_bdim] = unwrapTensorAtLevel(value, cur_level);
  25128. std::optional<Tensor> cum_seq_q_value;
  25129. std::optional<int64_t> cum_seq_q_bdim;
  25130. if (cum_seq_q) {
  25131. std::tie(cum_seq_q_value, cum_seq_q_bdim) = unwrapTensorAtLevel(cum_seq_q.value(), cur_level);
  25132. }
  25133. std::optional<Tensor> cum_seq_k_value;
  25134. std::optional<int64_t> cum_seq_k_bdim;
  25135. if (cum_seq_k) {
  25136. std::tie(cum_seq_k_value, cum_seq_k_bdim) = unwrapTensorAtLevel(cum_seq_k.value(), cur_level);
  25137. }
  25138. std::optional<Tensor> seqused_k_value;
  25139. std::optional<int64_t> seqused_k_bdim;
  25140. if (seqused_k) {
  25141. std::tie(seqused_k_value, seqused_k_bdim) = unwrapTensorAtLevel(seqused_k.value(), cur_level);
  25142. }
  25143. std::optional<Tensor> alibi_slopes_value;
  25144. std::optional<int64_t> alibi_slopes_bdim;
  25145. if (alibi_slopes) {
  25146. std::tie(alibi_slopes_value, alibi_slopes_bdim) = unwrapTensorAtLevel(alibi_slopes.value(), cur_level);
  25147. }
  25148. auto results = batch_rule(query_value, query_bdim, key_value, key_bdim, value_value, value_bdim, cum_seq_q_value, cum_seq_q_bdim, cum_seq_k_value, cum_seq_k_bdim, max_q, max_k, dropout_p, is_causal, return_debug_mask, scale, window_size_left, window_size_right, seqused_k_value, seqused_k_bdim, alibi_slopes_value, alibi_slopes_bdim);
  25149. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level), makeBatched(std::get<6>(results), std::get<7>(results), cur_level), makeBatched(std::get<8>(results), std::get<9>(results), cur_level));
  25150. }
  25151. template <typename batch_rule_t, batch_rule_t batch_rule>
  25152. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> _flash_attention_backward_generated_plumbing(const at::Tensor & grad_out, const at::Tensor & query, const at::Tensor & key, const at::Tensor & value, const at::Tensor & out, const at::Tensor & logsumexp, const at::Tensor & cum_seq_q, const at::Tensor & cum_seq_k, c10::SymInt max_q, c10::SymInt max_k, double dropout_p, bool is_causal, const at::Tensor & rng_state, const at::Tensor & unused, ::std::optional<double> scale, ::std::optional<c10::SymInt> window_size_left, ::std::optional<c10::SymInt> window_size_right) {
  25153. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25154. auto maybe_layer = maybeCurrentDynamicLayer();
  25155. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25156. int64_t cur_level = maybe_layer->layerId();
  25157. if (!isBatchedAtLevel(grad_out, cur_level) && !isBatchedAtLevel(query, cur_level) && !isBatchedAtLevel(key, cur_level) && !isBatchedAtLevel(value, cur_level) && !isBatchedAtLevel(out, cur_level) && !isBatchedAtLevel(logsumexp, cur_level) && !isBatchedAtLevel(cum_seq_q, cur_level) && !isBatchedAtLevel(cum_seq_k, cur_level) && !isBatchedAtLevel(rng_state, cur_level) && !isBatchedAtLevel(unused, cur_level)) {
  25158. return at::_ops::_flash_attention_backward::call(grad_out, query, key, value, out, logsumexp, cum_seq_q, cum_seq_k, max_q, max_k, dropout_p, is_causal, rng_state, unused, scale, window_size_left, window_size_right);
  25159. }
  25160. auto [grad_out_value, grad_out_bdim] = unwrapTensorAtLevel(grad_out, cur_level);
  25161. auto [query_value, query_bdim] = unwrapTensorAtLevel(query, cur_level);
  25162. auto [key_value, key_bdim] = unwrapTensorAtLevel(key, cur_level);
  25163. auto [value_value, value_bdim] = unwrapTensorAtLevel(value, cur_level);
  25164. auto [out_value, out_bdim] = unwrapTensorAtLevel(out, cur_level);
  25165. auto [logsumexp_value, logsumexp_bdim] = unwrapTensorAtLevel(logsumexp, cur_level);
  25166. auto [cum_seq_q_value, cum_seq_q_bdim] = unwrapTensorAtLevel(cum_seq_q, cur_level);
  25167. auto [cum_seq_k_value, cum_seq_k_bdim] = unwrapTensorAtLevel(cum_seq_k, cur_level);
  25168. auto [rng_state_value, rng_state_bdim] = unwrapTensorAtLevel(rng_state, cur_level);
  25169. auto [unused_value, unused_bdim] = unwrapTensorAtLevel(unused, cur_level);
  25170. auto results = batch_rule(grad_out_value, grad_out_bdim, query_value, query_bdim, key_value, key_bdim, value_value, value_bdim, out_value, out_bdim, logsumexp_value, logsumexp_bdim, cum_seq_q_value, cum_seq_q_bdim, cum_seq_k_value, cum_seq_k_bdim, max_q, max_k, dropout_p, is_causal, rng_state_value, rng_state_bdim, unused_value, unused_bdim, scale, window_size_left, window_size_right);
  25171. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  25172. }
  25173. template <typename batch_rule_t, batch_rule_t batch_rule>
  25174. ::std::tuple<at::Tensor,at::Tensor,at::Tensor,at::Tensor> _efficient_attention_backward_generated_plumbing(const at::Tensor & grad_out_, const at::Tensor & query, const at::Tensor & key, const at::Tensor & value, const ::std::optional<at::Tensor> & bias, const at::Tensor & out, const ::std::optional<at::Tensor> & cu_seqlens_q, const ::std::optional<at::Tensor> & cu_seqlens_k, c10::SymInt max_seqlen_q, c10::SymInt max_seqlen_k, const at::Tensor & logsumexp, double dropout_p, const at::Tensor & philox_seed, const at::Tensor & philox_offset, int64_t custom_mask_type, bool bias_requires_grad, ::std::optional<double> scale, ::std::optional<int64_t> num_splits_key, ::std::optional<int64_t> window_size, bool shared_storage_dqdkdv) {
  25175. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25176. auto maybe_layer = maybeCurrentDynamicLayer();
  25177. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25178. int64_t cur_level = maybe_layer->layerId();
  25179. if (!isBatchedAtLevel(grad_out_, cur_level) && !isBatchedAtLevel(query, cur_level) && !isBatchedAtLevel(key, cur_level) && !isBatchedAtLevel(value, cur_level) && !isBatchedAtLevel(bias, cur_level) && !isBatchedAtLevel(out, cur_level) && !isBatchedAtLevel(cu_seqlens_q, cur_level) && !isBatchedAtLevel(cu_seqlens_k, cur_level) && !isBatchedAtLevel(logsumexp, cur_level) && !isBatchedAtLevel(philox_seed, cur_level) && !isBatchedAtLevel(philox_offset, cur_level)) {
  25180. return at::_ops::_efficient_attention_backward::call(grad_out_, query, key, value, bias, out, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, logsumexp, dropout_p, philox_seed, philox_offset, custom_mask_type, bias_requires_grad, scale, num_splits_key, window_size, shared_storage_dqdkdv);
  25181. }
  25182. auto [grad_out__value, grad_out__bdim] = unwrapTensorAtLevel(grad_out_, cur_level);
  25183. auto [query_value, query_bdim] = unwrapTensorAtLevel(query, cur_level);
  25184. auto [key_value, key_bdim] = unwrapTensorAtLevel(key, cur_level);
  25185. auto [value_value, value_bdim] = unwrapTensorAtLevel(value, cur_level);
  25186. auto [out_value, out_bdim] = unwrapTensorAtLevel(out, cur_level);
  25187. auto [logsumexp_value, logsumexp_bdim] = unwrapTensorAtLevel(logsumexp, cur_level);
  25188. auto [philox_seed_value, philox_seed_bdim] = unwrapTensorAtLevel(philox_seed, cur_level);
  25189. auto [philox_offset_value, philox_offset_bdim] = unwrapTensorAtLevel(philox_offset, cur_level);
  25190. std::optional<Tensor> bias_value;
  25191. std::optional<int64_t> bias_bdim;
  25192. if (bias) {
  25193. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  25194. }
  25195. std::optional<Tensor> cu_seqlens_q_value;
  25196. std::optional<int64_t> cu_seqlens_q_bdim;
  25197. if (cu_seqlens_q) {
  25198. std::tie(cu_seqlens_q_value, cu_seqlens_q_bdim) = unwrapTensorAtLevel(cu_seqlens_q.value(), cur_level);
  25199. }
  25200. std::optional<Tensor> cu_seqlens_k_value;
  25201. std::optional<int64_t> cu_seqlens_k_bdim;
  25202. if (cu_seqlens_k) {
  25203. std::tie(cu_seqlens_k_value, cu_seqlens_k_bdim) = unwrapTensorAtLevel(cu_seqlens_k.value(), cur_level);
  25204. }
  25205. auto results = batch_rule(grad_out__value, grad_out__bdim, query_value, query_bdim, key_value, key_bdim, value_value, value_bdim, bias_value, bias_bdim, out_value, out_bdim, cu_seqlens_q_value, cu_seqlens_q_bdim, cu_seqlens_k_value, cu_seqlens_k_bdim, max_seqlen_q, max_seqlen_k, logsumexp_value, logsumexp_bdim, dropout_p, philox_seed_value, philox_seed_bdim, philox_offset_value, philox_offset_bdim, custom_mask_type, bias_requires_grad, scale, num_splits_key, window_size, shared_storage_dqdkdv);
  25206. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level), makeBatched(std::get<6>(results), std::get<7>(results), cur_level));
  25207. }
  25208. template <typename batch_rule_t, batch_rule_t batch_rule>
  25209. ::std::tuple<at::Tensor,at::Tensor,at::Tensor> _cudnn_attention_backward_generated_plumbing(const at::Tensor & grad_out, const at::Tensor & query, const at::Tensor & key, const at::Tensor & value, const at::Tensor & out, const at::Tensor & logsumexp, const at::Tensor & philox_seed, const at::Tensor & philox_offset, const at::Tensor & attn_bias, const at::Tensor & cum_seq_q, const at::Tensor & cum_seq_k, c10::SymInt max_q, c10::SymInt max_k, double dropout_p, bool is_causal, ::std::optional<double> scale) {
  25210. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25211. auto maybe_layer = maybeCurrentDynamicLayer();
  25212. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25213. int64_t cur_level = maybe_layer->layerId();
  25214. if (!isBatchedAtLevel(grad_out, cur_level) && !isBatchedAtLevel(query, cur_level) && !isBatchedAtLevel(key, cur_level) && !isBatchedAtLevel(value, cur_level) && !isBatchedAtLevel(out, cur_level) && !isBatchedAtLevel(logsumexp, cur_level) && !isBatchedAtLevel(philox_seed, cur_level) && !isBatchedAtLevel(philox_offset, cur_level) && !isBatchedAtLevel(attn_bias, cur_level) && !isBatchedAtLevel(cum_seq_q, cur_level) && !isBatchedAtLevel(cum_seq_k, cur_level)) {
  25215. return at::_ops::_cudnn_attention_backward::call(grad_out, query, key, value, out, logsumexp, philox_seed, philox_offset, attn_bias, cum_seq_q, cum_seq_k, max_q, max_k, dropout_p, is_causal, scale);
  25216. }
  25217. auto [grad_out_value, grad_out_bdim] = unwrapTensorAtLevel(grad_out, cur_level);
  25218. auto [query_value, query_bdim] = unwrapTensorAtLevel(query, cur_level);
  25219. auto [key_value, key_bdim] = unwrapTensorAtLevel(key, cur_level);
  25220. auto [value_value, value_bdim] = unwrapTensorAtLevel(value, cur_level);
  25221. auto [out_value, out_bdim] = unwrapTensorAtLevel(out, cur_level);
  25222. auto [logsumexp_value, logsumexp_bdim] = unwrapTensorAtLevel(logsumexp, cur_level);
  25223. auto [philox_seed_value, philox_seed_bdim] = unwrapTensorAtLevel(philox_seed, cur_level);
  25224. auto [philox_offset_value, philox_offset_bdim] = unwrapTensorAtLevel(philox_offset, cur_level);
  25225. auto [attn_bias_value, attn_bias_bdim] = unwrapTensorAtLevel(attn_bias, cur_level);
  25226. auto [cum_seq_q_value, cum_seq_q_bdim] = unwrapTensorAtLevel(cum_seq_q, cur_level);
  25227. auto [cum_seq_k_value, cum_seq_k_bdim] = unwrapTensorAtLevel(cum_seq_k, cur_level);
  25228. auto results = batch_rule(grad_out_value, grad_out_bdim, query_value, query_bdim, key_value, key_bdim, value_value, value_bdim, out_value, out_bdim, logsumexp_value, logsumexp_bdim, philox_seed_value, philox_seed_bdim, philox_offset_value, philox_offset_bdim, attn_bias_value, attn_bias_bdim, cum_seq_q_value, cum_seq_q_bdim, cum_seq_k_value, cum_seq_k_bdim, max_q, max_k, dropout_p, is_causal, scale);
  25229. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level));
  25230. }
  25231. template <typename batch_rule_t, batch_rule_t batch_rule>
  25232. at::Tensor _triton_scaled_dot_attention_generated_plumbing(const at::Tensor & q, const at::Tensor & k, const at::Tensor & v, double dropout_p) {
  25233. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25234. auto maybe_layer = maybeCurrentDynamicLayer();
  25235. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25236. int64_t cur_level = maybe_layer->layerId();
  25237. if (!isBatchedAtLevel(q, cur_level) && !isBatchedAtLevel(k, cur_level) && !isBatchedAtLevel(v, cur_level)) {
  25238. return at::_ops::_triton_scaled_dot_attention::call(q, k, v, dropout_p);
  25239. }
  25240. auto [q_value, q_bdim] = unwrapTensorAtLevel(q, cur_level);
  25241. auto [k_value, k_bdim] = unwrapTensorAtLevel(k, cur_level);
  25242. auto [v_value, v_bdim] = unwrapTensorAtLevel(v, cur_level);
  25243. auto results = batch_rule(q_value, q_bdim, k_value, k_bdim, v_value, v_bdim, dropout_p);
  25244. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25245. }
  25246. template <typename batch_rule_t, batch_rule_t batch_rule>
  25247. at::Tensor & _fill_mem_eff_dropout_mask__generated_plumbing(at::Tensor & self, double dropout_p, int64_t seed, int64_t offset) {
  25248. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25249. auto maybe_layer = maybeCurrentDynamicLayer();
  25250. vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing");
  25251. int64_t cur_level = maybe_layer->layerId();
  25252. if (!isBatchedAtLevel(self, cur_level)) {
  25253. return at::_ops::_fill_mem_eff_dropout_mask_::call(self, dropout_p, seed, offset);
  25254. }
  25255. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  25256. batch_rule(self_value, self_bdim, dropout_p, seed, offset);
  25257. return self;
  25258. }
  25259. template <typename batch_rule_t, batch_rule_t batch_rule>
  25260. at::Tensor _triton_multi_head_attention_generated_plumbing(const at::Tensor & query, const at::Tensor & key, const at::Tensor & value, int64_t embed_dim, int64_t num_head, const at::Tensor & qkv_weight, const at::Tensor & qkv_bias, const at::Tensor & proj_weight, const at::Tensor & proj_bias, const ::std::optional<at::Tensor> & mask) {
  25261. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25262. auto maybe_layer = maybeCurrentDynamicLayer();
  25263. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25264. int64_t cur_level = maybe_layer->layerId();
  25265. if (!isBatchedAtLevel(query, cur_level) && !isBatchedAtLevel(key, cur_level) && !isBatchedAtLevel(value, cur_level) && !isBatchedAtLevel(qkv_weight, cur_level) && !isBatchedAtLevel(qkv_bias, cur_level) && !isBatchedAtLevel(proj_weight, cur_level) && !isBatchedAtLevel(proj_bias, cur_level) && !isBatchedAtLevel(mask, cur_level)) {
  25266. return at::_ops::_triton_multi_head_attention::call(query, key, value, embed_dim, num_head, qkv_weight, qkv_bias, proj_weight, proj_bias, mask);
  25267. }
  25268. auto [query_value, query_bdim] = unwrapTensorAtLevel(query, cur_level);
  25269. auto [key_value, key_bdim] = unwrapTensorAtLevel(key, cur_level);
  25270. auto [value_value, value_bdim] = unwrapTensorAtLevel(value, cur_level);
  25271. auto [qkv_weight_value, qkv_weight_bdim] = unwrapTensorAtLevel(qkv_weight, cur_level);
  25272. auto [qkv_bias_value, qkv_bias_bdim] = unwrapTensorAtLevel(qkv_bias, cur_level);
  25273. auto [proj_weight_value, proj_weight_bdim] = unwrapTensorAtLevel(proj_weight, cur_level);
  25274. auto [proj_bias_value, proj_bias_bdim] = unwrapTensorAtLevel(proj_bias, cur_level);
  25275. std::optional<Tensor> mask_value;
  25276. std::optional<int64_t> mask_bdim;
  25277. if (mask) {
  25278. std::tie(mask_value, mask_bdim) = unwrapTensorAtLevel(mask.value(), cur_level);
  25279. }
  25280. auto results = batch_rule(query_value, query_bdim, key_value, key_bdim, value_value, value_bdim, embed_dim, num_head, qkv_weight_value, qkv_weight_bdim, qkv_bias_value, qkv_bias_bdim, proj_weight_value, proj_weight_bdim, proj_bias_value, proj_bias_bdim, mask_value, mask_bdim);
  25281. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25282. }
  25283. template <typename batch_rule_t, batch_rule_t batch_rule>
  25284. at::Tensor special_airy_ai_generated_plumbing(const at::Tensor & x) {
  25285. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25286. auto maybe_layer = maybeCurrentDynamicLayer();
  25287. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25288. int64_t cur_level = maybe_layer->layerId();
  25289. if (!isBatchedAtLevel(x, cur_level)) {
  25290. return at::_ops::special_airy_ai::call(x);
  25291. }
  25292. auto [x_value, x_bdim] = unwrapTensorAtLevel(x, cur_level);
  25293. auto results = batch_rule(x_value, x_bdim);
  25294. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25295. }
  25296. template <typename batch_rule_t, batch_rule_t batch_rule>
  25297. at::Tensor special_bessel_j0_generated_plumbing(const at::Tensor & self) {
  25298. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25299. auto maybe_layer = maybeCurrentDynamicLayer();
  25300. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25301. int64_t cur_level = maybe_layer->layerId();
  25302. if (!isBatchedAtLevel(self, cur_level)) {
  25303. return at::_ops::special_bessel_j0::call(self);
  25304. }
  25305. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  25306. auto results = batch_rule(self_value, self_bdim);
  25307. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25308. }
  25309. template <typename batch_rule_t, batch_rule_t batch_rule>
  25310. at::Tensor special_bessel_j1_generated_plumbing(const at::Tensor & self) {
  25311. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25312. auto maybe_layer = maybeCurrentDynamicLayer();
  25313. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25314. int64_t cur_level = maybe_layer->layerId();
  25315. if (!isBatchedAtLevel(self, cur_level)) {
  25316. return at::_ops::special_bessel_j1::call(self);
  25317. }
  25318. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  25319. auto results = batch_rule(self_value, self_bdim);
  25320. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25321. }
  25322. template <typename batch_rule_t, batch_rule_t batch_rule>
  25323. at::Tensor special_bessel_y0_generated_plumbing(const at::Tensor & self) {
  25324. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25325. auto maybe_layer = maybeCurrentDynamicLayer();
  25326. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25327. int64_t cur_level = maybe_layer->layerId();
  25328. if (!isBatchedAtLevel(self, cur_level)) {
  25329. return at::_ops::special_bessel_y0::call(self);
  25330. }
  25331. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  25332. auto results = batch_rule(self_value, self_bdim);
  25333. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25334. }
  25335. template <typename batch_rule_t, batch_rule_t batch_rule>
  25336. at::Tensor special_bessel_y1_generated_plumbing(const at::Tensor & self) {
  25337. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25338. auto maybe_layer = maybeCurrentDynamicLayer();
  25339. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25340. int64_t cur_level = maybe_layer->layerId();
  25341. if (!isBatchedAtLevel(self, cur_level)) {
  25342. return at::_ops::special_bessel_y1::call(self);
  25343. }
  25344. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  25345. auto results = batch_rule(self_value, self_bdim);
  25346. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25347. }
  25348. template <typename batch_rule_t, batch_rule_t batch_rule>
  25349. at::Tensor special_chebyshev_polynomial_t_generated_plumbing(const at::Tensor & x, const at::Tensor & n) {
  25350. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25351. auto maybe_layer = maybeCurrentDynamicLayer();
  25352. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25353. int64_t cur_level = maybe_layer->layerId();
  25354. if (!isBatchedAtLevel(x, cur_level) && !isBatchedAtLevel(n, cur_level)) {
  25355. return at::_ops::special_chebyshev_polynomial_t::call(x, n);
  25356. }
  25357. auto [x_value, x_bdim] = unwrapTensorAtLevel(x, cur_level);
  25358. auto [n_value, n_bdim] = unwrapTensorAtLevel(n, cur_level);
  25359. auto results = batch_rule(x_value, x_bdim, n_value, n_bdim);
  25360. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25361. }
  25362. template <typename batch_rule_t, batch_rule_t batch_rule>
  25363. at::Tensor special_chebyshev_polynomial_t_x_scalar_generated_plumbing(const at::Scalar & x, const at::Tensor & n) {
  25364. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25365. auto maybe_layer = maybeCurrentDynamicLayer();
  25366. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25367. int64_t cur_level = maybe_layer->layerId();
  25368. if (!isBatchedAtLevel(n, cur_level)) {
  25369. return at::_ops::special_chebyshev_polynomial_t_x_scalar::call(x, n);
  25370. }
  25371. auto [n_value, n_bdim] = unwrapTensorAtLevel(n, cur_level);
  25372. auto results = batch_rule(x, n_value, n_bdim);
  25373. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25374. }
  25375. template <typename batch_rule_t, batch_rule_t batch_rule>
  25376. at::Tensor special_chebyshev_polynomial_t_n_scalar_generated_plumbing(const at::Tensor & x, const at::Scalar & n) {
  25377. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25378. auto maybe_layer = maybeCurrentDynamicLayer();
  25379. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25380. int64_t cur_level = maybe_layer->layerId();
  25381. if (!isBatchedAtLevel(x, cur_level)) {
  25382. return at::_ops::special_chebyshev_polynomial_t_n_scalar::call(x, n);
  25383. }
  25384. auto [x_value, x_bdim] = unwrapTensorAtLevel(x, cur_level);
  25385. auto results = batch_rule(x_value, x_bdim, n);
  25386. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25387. }
  25388. template <typename batch_rule_t, batch_rule_t batch_rule>
  25389. at::Tensor special_chebyshev_polynomial_u_generated_plumbing(const at::Tensor & x, const at::Tensor & n) {
  25390. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25391. auto maybe_layer = maybeCurrentDynamicLayer();
  25392. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25393. int64_t cur_level = maybe_layer->layerId();
  25394. if (!isBatchedAtLevel(x, cur_level) && !isBatchedAtLevel(n, cur_level)) {
  25395. return at::_ops::special_chebyshev_polynomial_u::call(x, n);
  25396. }
  25397. auto [x_value, x_bdim] = unwrapTensorAtLevel(x, cur_level);
  25398. auto [n_value, n_bdim] = unwrapTensorAtLevel(n, cur_level);
  25399. auto results = batch_rule(x_value, x_bdim, n_value, n_bdim);
  25400. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25401. }
  25402. template <typename batch_rule_t, batch_rule_t batch_rule>
  25403. at::Tensor special_chebyshev_polynomial_u_x_scalar_generated_plumbing(const at::Scalar & x, const at::Tensor & n) {
  25404. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25405. auto maybe_layer = maybeCurrentDynamicLayer();
  25406. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25407. int64_t cur_level = maybe_layer->layerId();
  25408. if (!isBatchedAtLevel(n, cur_level)) {
  25409. return at::_ops::special_chebyshev_polynomial_u_x_scalar::call(x, n);
  25410. }
  25411. auto [n_value, n_bdim] = unwrapTensorAtLevel(n, cur_level);
  25412. auto results = batch_rule(x, n_value, n_bdim);
  25413. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25414. }
  25415. template <typename batch_rule_t, batch_rule_t batch_rule>
  25416. at::Tensor special_chebyshev_polynomial_u_n_scalar_generated_plumbing(const at::Tensor & x, const at::Scalar & n) {
  25417. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25418. auto maybe_layer = maybeCurrentDynamicLayer();
  25419. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25420. int64_t cur_level = maybe_layer->layerId();
  25421. if (!isBatchedAtLevel(x, cur_level)) {
  25422. return at::_ops::special_chebyshev_polynomial_u_n_scalar::call(x, n);
  25423. }
  25424. auto [x_value, x_bdim] = unwrapTensorAtLevel(x, cur_level);
  25425. auto results = batch_rule(x_value, x_bdim, n);
  25426. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25427. }
  25428. template <typename batch_rule_t, batch_rule_t batch_rule>
  25429. at::Tensor special_chebyshev_polynomial_v_generated_plumbing(const at::Tensor & x, const at::Tensor & n) {
  25430. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25431. auto maybe_layer = maybeCurrentDynamicLayer();
  25432. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25433. int64_t cur_level = maybe_layer->layerId();
  25434. if (!isBatchedAtLevel(x, cur_level) && !isBatchedAtLevel(n, cur_level)) {
  25435. return at::_ops::special_chebyshev_polynomial_v::call(x, n);
  25436. }
  25437. auto [x_value, x_bdim] = unwrapTensorAtLevel(x, cur_level);
  25438. auto [n_value, n_bdim] = unwrapTensorAtLevel(n, cur_level);
  25439. auto results = batch_rule(x_value, x_bdim, n_value, n_bdim);
  25440. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25441. }
  25442. template <typename batch_rule_t, batch_rule_t batch_rule>
  25443. at::Tensor special_chebyshev_polynomial_v_x_scalar_generated_plumbing(const at::Scalar & x, const at::Tensor & n) {
  25444. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25445. auto maybe_layer = maybeCurrentDynamicLayer();
  25446. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25447. int64_t cur_level = maybe_layer->layerId();
  25448. if (!isBatchedAtLevel(n, cur_level)) {
  25449. return at::_ops::special_chebyshev_polynomial_v_x_scalar::call(x, n);
  25450. }
  25451. auto [n_value, n_bdim] = unwrapTensorAtLevel(n, cur_level);
  25452. auto results = batch_rule(x, n_value, n_bdim);
  25453. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25454. }
  25455. template <typename batch_rule_t, batch_rule_t batch_rule>
  25456. at::Tensor special_chebyshev_polynomial_v_n_scalar_generated_plumbing(const at::Tensor & x, const at::Scalar & n) {
  25457. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25458. auto maybe_layer = maybeCurrentDynamicLayer();
  25459. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25460. int64_t cur_level = maybe_layer->layerId();
  25461. if (!isBatchedAtLevel(x, cur_level)) {
  25462. return at::_ops::special_chebyshev_polynomial_v_n_scalar::call(x, n);
  25463. }
  25464. auto [x_value, x_bdim] = unwrapTensorAtLevel(x, cur_level);
  25465. auto results = batch_rule(x_value, x_bdim, n);
  25466. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25467. }
  25468. template <typename batch_rule_t, batch_rule_t batch_rule>
  25469. at::Tensor special_chebyshev_polynomial_w_generated_plumbing(const at::Tensor & x, const at::Tensor & n) {
  25470. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25471. auto maybe_layer = maybeCurrentDynamicLayer();
  25472. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25473. int64_t cur_level = maybe_layer->layerId();
  25474. if (!isBatchedAtLevel(x, cur_level) && !isBatchedAtLevel(n, cur_level)) {
  25475. return at::_ops::special_chebyshev_polynomial_w::call(x, n);
  25476. }
  25477. auto [x_value, x_bdim] = unwrapTensorAtLevel(x, cur_level);
  25478. auto [n_value, n_bdim] = unwrapTensorAtLevel(n, cur_level);
  25479. auto results = batch_rule(x_value, x_bdim, n_value, n_bdim);
  25480. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25481. }
  25482. template <typename batch_rule_t, batch_rule_t batch_rule>
  25483. at::Tensor special_chebyshev_polynomial_w_x_scalar_generated_plumbing(const at::Scalar & x, const at::Tensor & n) {
  25484. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25485. auto maybe_layer = maybeCurrentDynamicLayer();
  25486. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25487. int64_t cur_level = maybe_layer->layerId();
  25488. if (!isBatchedAtLevel(n, cur_level)) {
  25489. return at::_ops::special_chebyshev_polynomial_w_x_scalar::call(x, n);
  25490. }
  25491. auto [n_value, n_bdim] = unwrapTensorAtLevel(n, cur_level);
  25492. auto results = batch_rule(x, n_value, n_bdim);
  25493. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25494. }
  25495. template <typename batch_rule_t, batch_rule_t batch_rule>
  25496. at::Tensor special_chebyshev_polynomial_w_n_scalar_generated_plumbing(const at::Tensor & x, const at::Scalar & n) {
  25497. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25498. auto maybe_layer = maybeCurrentDynamicLayer();
  25499. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25500. int64_t cur_level = maybe_layer->layerId();
  25501. if (!isBatchedAtLevel(x, cur_level)) {
  25502. return at::_ops::special_chebyshev_polynomial_w_n_scalar::call(x, n);
  25503. }
  25504. auto [x_value, x_bdim] = unwrapTensorAtLevel(x, cur_level);
  25505. auto results = batch_rule(x_value, x_bdim, n);
  25506. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25507. }
  25508. template <typename batch_rule_t, batch_rule_t batch_rule>
  25509. at::Tensor special_hermite_polynomial_h_generated_plumbing(const at::Tensor & x, const at::Tensor & n) {
  25510. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25511. auto maybe_layer = maybeCurrentDynamicLayer();
  25512. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25513. int64_t cur_level = maybe_layer->layerId();
  25514. if (!isBatchedAtLevel(x, cur_level) && !isBatchedAtLevel(n, cur_level)) {
  25515. return at::_ops::special_hermite_polynomial_h::call(x, n);
  25516. }
  25517. auto [x_value, x_bdim] = unwrapTensorAtLevel(x, cur_level);
  25518. auto [n_value, n_bdim] = unwrapTensorAtLevel(n, cur_level);
  25519. auto results = batch_rule(x_value, x_bdim, n_value, n_bdim);
  25520. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25521. }
  25522. template <typename batch_rule_t, batch_rule_t batch_rule>
  25523. at::Tensor special_hermite_polynomial_h_x_scalar_generated_plumbing(const at::Scalar & x, const at::Tensor & n) {
  25524. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25525. auto maybe_layer = maybeCurrentDynamicLayer();
  25526. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25527. int64_t cur_level = maybe_layer->layerId();
  25528. if (!isBatchedAtLevel(n, cur_level)) {
  25529. return at::_ops::special_hermite_polynomial_h_x_scalar::call(x, n);
  25530. }
  25531. auto [n_value, n_bdim] = unwrapTensorAtLevel(n, cur_level);
  25532. auto results = batch_rule(x, n_value, n_bdim);
  25533. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25534. }
  25535. template <typename batch_rule_t, batch_rule_t batch_rule>
  25536. at::Tensor special_hermite_polynomial_h_n_scalar_generated_plumbing(const at::Tensor & x, const at::Scalar & n) {
  25537. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25538. auto maybe_layer = maybeCurrentDynamicLayer();
  25539. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25540. int64_t cur_level = maybe_layer->layerId();
  25541. if (!isBatchedAtLevel(x, cur_level)) {
  25542. return at::_ops::special_hermite_polynomial_h_n_scalar::call(x, n);
  25543. }
  25544. auto [x_value, x_bdim] = unwrapTensorAtLevel(x, cur_level);
  25545. auto results = batch_rule(x_value, x_bdim, n);
  25546. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25547. }
  25548. template <typename batch_rule_t, batch_rule_t batch_rule>
  25549. at::Tensor special_hermite_polynomial_he_generated_plumbing(const at::Tensor & x, const at::Tensor & n) {
  25550. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25551. auto maybe_layer = maybeCurrentDynamicLayer();
  25552. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25553. int64_t cur_level = maybe_layer->layerId();
  25554. if (!isBatchedAtLevel(x, cur_level) && !isBatchedAtLevel(n, cur_level)) {
  25555. return at::_ops::special_hermite_polynomial_he::call(x, n);
  25556. }
  25557. auto [x_value, x_bdim] = unwrapTensorAtLevel(x, cur_level);
  25558. auto [n_value, n_bdim] = unwrapTensorAtLevel(n, cur_level);
  25559. auto results = batch_rule(x_value, x_bdim, n_value, n_bdim);
  25560. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25561. }
  25562. template <typename batch_rule_t, batch_rule_t batch_rule>
  25563. at::Tensor special_hermite_polynomial_he_x_scalar_generated_plumbing(const at::Scalar & x, const at::Tensor & n) {
  25564. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25565. auto maybe_layer = maybeCurrentDynamicLayer();
  25566. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25567. int64_t cur_level = maybe_layer->layerId();
  25568. if (!isBatchedAtLevel(n, cur_level)) {
  25569. return at::_ops::special_hermite_polynomial_he_x_scalar::call(x, n);
  25570. }
  25571. auto [n_value, n_bdim] = unwrapTensorAtLevel(n, cur_level);
  25572. auto results = batch_rule(x, n_value, n_bdim);
  25573. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25574. }
  25575. template <typename batch_rule_t, batch_rule_t batch_rule>
  25576. at::Tensor special_hermite_polynomial_he_n_scalar_generated_plumbing(const at::Tensor & x, const at::Scalar & n) {
  25577. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25578. auto maybe_layer = maybeCurrentDynamicLayer();
  25579. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25580. int64_t cur_level = maybe_layer->layerId();
  25581. if (!isBatchedAtLevel(x, cur_level)) {
  25582. return at::_ops::special_hermite_polynomial_he_n_scalar::call(x, n);
  25583. }
  25584. auto [x_value, x_bdim] = unwrapTensorAtLevel(x, cur_level);
  25585. auto results = batch_rule(x_value, x_bdim, n);
  25586. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25587. }
  25588. template <typename batch_rule_t, batch_rule_t batch_rule>
  25589. at::Tensor special_laguerre_polynomial_l_generated_plumbing(const at::Tensor & x, const at::Tensor & n) {
  25590. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25591. auto maybe_layer = maybeCurrentDynamicLayer();
  25592. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25593. int64_t cur_level = maybe_layer->layerId();
  25594. if (!isBatchedAtLevel(x, cur_level) && !isBatchedAtLevel(n, cur_level)) {
  25595. return at::_ops::special_laguerre_polynomial_l::call(x, n);
  25596. }
  25597. auto [x_value, x_bdim] = unwrapTensorAtLevel(x, cur_level);
  25598. auto [n_value, n_bdim] = unwrapTensorAtLevel(n, cur_level);
  25599. auto results = batch_rule(x_value, x_bdim, n_value, n_bdim);
  25600. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25601. }
  25602. template <typename batch_rule_t, batch_rule_t batch_rule>
  25603. at::Tensor special_laguerre_polynomial_l_x_scalar_generated_plumbing(const at::Scalar & x, const at::Tensor & n) {
  25604. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25605. auto maybe_layer = maybeCurrentDynamicLayer();
  25606. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25607. int64_t cur_level = maybe_layer->layerId();
  25608. if (!isBatchedAtLevel(n, cur_level)) {
  25609. return at::_ops::special_laguerre_polynomial_l_x_scalar::call(x, n);
  25610. }
  25611. auto [n_value, n_bdim] = unwrapTensorAtLevel(n, cur_level);
  25612. auto results = batch_rule(x, n_value, n_bdim);
  25613. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25614. }
  25615. template <typename batch_rule_t, batch_rule_t batch_rule>
  25616. at::Tensor special_laguerre_polynomial_l_n_scalar_generated_plumbing(const at::Tensor & x, const at::Scalar & n) {
  25617. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25618. auto maybe_layer = maybeCurrentDynamicLayer();
  25619. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25620. int64_t cur_level = maybe_layer->layerId();
  25621. if (!isBatchedAtLevel(x, cur_level)) {
  25622. return at::_ops::special_laguerre_polynomial_l_n_scalar::call(x, n);
  25623. }
  25624. auto [x_value, x_bdim] = unwrapTensorAtLevel(x, cur_level);
  25625. auto results = batch_rule(x_value, x_bdim, n);
  25626. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25627. }
  25628. template <typename batch_rule_t, batch_rule_t batch_rule>
  25629. at::Tensor special_legendre_polynomial_p_generated_plumbing(const at::Tensor & x, const at::Tensor & n) {
  25630. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25631. auto maybe_layer = maybeCurrentDynamicLayer();
  25632. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25633. int64_t cur_level = maybe_layer->layerId();
  25634. if (!isBatchedAtLevel(x, cur_level) && !isBatchedAtLevel(n, cur_level)) {
  25635. return at::_ops::special_legendre_polynomial_p::call(x, n);
  25636. }
  25637. auto [x_value, x_bdim] = unwrapTensorAtLevel(x, cur_level);
  25638. auto [n_value, n_bdim] = unwrapTensorAtLevel(n, cur_level);
  25639. auto results = batch_rule(x_value, x_bdim, n_value, n_bdim);
  25640. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25641. }
  25642. template <typename batch_rule_t, batch_rule_t batch_rule>
  25643. at::Tensor special_legendre_polynomial_p_x_scalar_generated_plumbing(const at::Scalar & x, const at::Tensor & n) {
  25644. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25645. auto maybe_layer = maybeCurrentDynamicLayer();
  25646. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25647. int64_t cur_level = maybe_layer->layerId();
  25648. if (!isBatchedAtLevel(n, cur_level)) {
  25649. return at::_ops::special_legendre_polynomial_p_x_scalar::call(x, n);
  25650. }
  25651. auto [n_value, n_bdim] = unwrapTensorAtLevel(n, cur_level);
  25652. auto results = batch_rule(x, n_value, n_bdim);
  25653. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25654. }
  25655. template <typename batch_rule_t, batch_rule_t batch_rule>
  25656. at::Tensor special_legendre_polynomial_p_n_scalar_generated_plumbing(const at::Tensor & x, const at::Scalar & n) {
  25657. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25658. auto maybe_layer = maybeCurrentDynamicLayer();
  25659. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25660. int64_t cur_level = maybe_layer->layerId();
  25661. if (!isBatchedAtLevel(x, cur_level)) {
  25662. return at::_ops::special_legendre_polynomial_p_n_scalar::call(x, n);
  25663. }
  25664. auto [x_value, x_bdim] = unwrapTensorAtLevel(x, cur_level);
  25665. auto results = batch_rule(x_value, x_bdim, n);
  25666. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25667. }
  25668. template <typename batch_rule_t, batch_rule_t batch_rule>
  25669. at::Tensor special_modified_bessel_i0_generated_plumbing(const at::Tensor & self) {
  25670. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25671. auto maybe_layer = maybeCurrentDynamicLayer();
  25672. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25673. int64_t cur_level = maybe_layer->layerId();
  25674. if (!isBatchedAtLevel(self, cur_level)) {
  25675. return at::_ops::special_modified_bessel_i0::call(self);
  25676. }
  25677. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  25678. auto results = batch_rule(self_value, self_bdim);
  25679. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25680. }
  25681. template <typename batch_rule_t, batch_rule_t batch_rule>
  25682. at::Tensor special_modified_bessel_i1_generated_plumbing(const at::Tensor & self) {
  25683. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25684. auto maybe_layer = maybeCurrentDynamicLayer();
  25685. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25686. int64_t cur_level = maybe_layer->layerId();
  25687. if (!isBatchedAtLevel(self, cur_level)) {
  25688. return at::_ops::special_modified_bessel_i1::call(self);
  25689. }
  25690. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  25691. auto results = batch_rule(self_value, self_bdim);
  25692. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25693. }
  25694. template <typename batch_rule_t, batch_rule_t batch_rule>
  25695. at::Tensor special_modified_bessel_k0_generated_plumbing(const at::Tensor & self) {
  25696. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25697. auto maybe_layer = maybeCurrentDynamicLayer();
  25698. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25699. int64_t cur_level = maybe_layer->layerId();
  25700. if (!isBatchedAtLevel(self, cur_level)) {
  25701. return at::_ops::special_modified_bessel_k0::call(self);
  25702. }
  25703. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  25704. auto results = batch_rule(self_value, self_bdim);
  25705. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25706. }
  25707. template <typename batch_rule_t, batch_rule_t batch_rule>
  25708. at::Tensor special_modified_bessel_k1_generated_plumbing(const at::Tensor & self) {
  25709. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25710. auto maybe_layer = maybeCurrentDynamicLayer();
  25711. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25712. int64_t cur_level = maybe_layer->layerId();
  25713. if (!isBatchedAtLevel(self, cur_level)) {
  25714. return at::_ops::special_modified_bessel_k1::call(self);
  25715. }
  25716. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  25717. auto results = batch_rule(self_value, self_bdim);
  25718. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25719. }
  25720. template <typename batch_rule_t, batch_rule_t batch_rule>
  25721. at::Tensor special_scaled_modified_bessel_k0_generated_plumbing(const at::Tensor & x) {
  25722. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25723. auto maybe_layer = maybeCurrentDynamicLayer();
  25724. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25725. int64_t cur_level = maybe_layer->layerId();
  25726. if (!isBatchedAtLevel(x, cur_level)) {
  25727. return at::_ops::special_scaled_modified_bessel_k0::call(x);
  25728. }
  25729. auto [x_value, x_bdim] = unwrapTensorAtLevel(x, cur_level);
  25730. auto results = batch_rule(x_value, x_bdim);
  25731. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25732. }
  25733. template <typename batch_rule_t, batch_rule_t batch_rule>
  25734. at::Tensor special_scaled_modified_bessel_k1_generated_plumbing(const at::Tensor & x) {
  25735. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25736. auto maybe_layer = maybeCurrentDynamicLayer();
  25737. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25738. int64_t cur_level = maybe_layer->layerId();
  25739. if (!isBatchedAtLevel(x, cur_level)) {
  25740. return at::_ops::special_scaled_modified_bessel_k1::call(x);
  25741. }
  25742. auto [x_value, x_bdim] = unwrapTensorAtLevel(x, cur_level);
  25743. auto results = batch_rule(x_value, x_bdim);
  25744. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25745. }
  25746. template <typename batch_rule_t, batch_rule_t batch_rule>
  25747. at::Tensor special_shifted_chebyshev_polynomial_t_generated_plumbing(const at::Tensor & x, const at::Tensor & n) {
  25748. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25749. auto maybe_layer = maybeCurrentDynamicLayer();
  25750. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25751. int64_t cur_level = maybe_layer->layerId();
  25752. if (!isBatchedAtLevel(x, cur_level) && !isBatchedAtLevel(n, cur_level)) {
  25753. return at::_ops::special_shifted_chebyshev_polynomial_t::call(x, n);
  25754. }
  25755. auto [x_value, x_bdim] = unwrapTensorAtLevel(x, cur_level);
  25756. auto [n_value, n_bdim] = unwrapTensorAtLevel(n, cur_level);
  25757. auto results = batch_rule(x_value, x_bdim, n_value, n_bdim);
  25758. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25759. }
  25760. template <typename batch_rule_t, batch_rule_t batch_rule>
  25761. at::Tensor special_shifted_chebyshev_polynomial_t_x_scalar_generated_plumbing(const at::Scalar & x, const at::Tensor & n) {
  25762. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25763. auto maybe_layer = maybeCurrentDynamicLayer();
  25764. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25765. int64_t cur_level = maybe_layer->layerId();
  25766. if (!isBatchedAtLevel(n, cur_level)) {
  25767. return at::_ops::special_shifted_chebyshev_polynomial_t_x_scalar::call(x, n);
  25768. }
  25769. auto [n_value, n_bdim] = unwrapTensorAtLevel(n, cur_level);
  25770. auto results = batch_rule(x, n_value, n_bdim);
  25771. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25772. }
  25773. template <typename batch_rule_t, batch_rule_t batch_rule>
  25774. at::Tensor special_shifted_chebyshev_polynomial_t_n_scalar_generated_plumbing(const at::Tensor & x, const at::Scalar & n) {
  25775. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25776. auto maybe_layer = maybeCurrentDynamicLayer();
  25777. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25778. int64_t cur_level = maybe_layer->layerId();
  25779. if (!isBatchedAtLevel(x, cur_level)) {
  25780. return at::_ops::special_shifted_chebyshev_polynomial_t_n_scalar::call(x, n);
  25781. }
  25782. auto [x_value, x_bdim] = unwrapTensorAtLevel(x, cur_level);
  25783. auto results = batch_rule(x_value, x_bdim, n);
  25784. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25785. }
  25786. template <typename batch_rule_t, batch_rule_t batch_rule>
  25787. at::Tensor special_shifted_chebyshev_polynomial_u_generated_plumbing(const at::Tensor & x, const at::Tensor & n) {
  25788. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25789. auto maybe_layer = maybeCurrentDynamicLayer();
  25790. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25791. int64_t cur_level = maybe_layer->layerId();
  25792. if (!isBatchedAtLevel(x, cur_level) && !isBatchedAtLevel(n, cur_level)) {
  25793. return at::_ops::special_shifted_chebyshev_polynomial_u::call(x, n);
  25794. }
  25795. auto [x_value, x_bdim] = unwrapTensorAtLevel(x, cur_level);
  25796. auto [n_value, n_bdim] = unwrapTensorAtLevel(n, cur_level);
  25797. auto results = batch_rule(x_value, x_bdim, n_value, n_bdim);
  25798. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25799. }
  25800. template <typename batch_rule_t, batch_rule_t batch_rule>
  25801. at::Tensor special_shifted_chebyshev_polynomial_u_x_scalar_generated_plumbing(const at::Scalar & x, const at::Tensor & n) {
  25802. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25803. auto maybe_layer = maybeCurrentDynamicLayer();
  25804. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25805. int64_t cur_level = maybe_layer->layerId();
  25806. if (!isBatchedAtLevel(n, cur_level)) {
  25807. return at::_ops::special_shifted_chebyshev_polynomial_u_x_scalar::call(x, n);
  25808. }
  25809. auto [n_value, n_bdim] = unwrapTensorAtLevel(n, cur_level);
  25810. auto results = batch_rule(x, n_value, n_bdim);
  25811. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25812. }
  25813. template <typename batch_rule_t, batch_rule_t batch_rule>
  25814. at::Tensor special_shifted_chebyshev_polynomial_u_n_scalar_generated_plumbing(const at::Tensor & x, const at::Scalar & n) {
  25815. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25816. auto maybe_layer = maybeCurrentDynamicLayer();
  25817. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25818. int64_t cur_level = maybe_layer->layerId();
  25819. if (!isBatchedAtLevel(x, cur_level)) {
  25820. return at::_ops::special_shifted_chebyshev_polynomial_u_n_scalar::call(x, n);
  25821. }
  25822. auto [x_value, x_bdim] = unwrapTensorAtLevel(x, cur_level);
  25823. auto results = batch_rule(x_value, x_bdim, n);
  25824. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25825. }
  25826. template <typename batch_rule_t, batch_rule_t batch_rule>
  25827. at::Tensor special_shifted_chebyshev_polynomial_v_generated_plumbing(const at::Tensor & x, const at::Tensor & n) {
  25828. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25829. auto maybe_layer = maybeCurrentDynamicLayer();
  25830. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25831. int64_t cur_level = maybe_layer->layerId();
  25832. if (!isBatchedAtLevel(x, cur_level) && !isBatchedAtLevel(n, cur_level)) {
  25833. return at::_ops::special_shifted_chebyshev_polynomial_v::call(x, n);
  25834. }
  25835. auto [x_value, x_bdim] = unwrapTensorAtLevel(x, cur_level);
  25836. auto [n_value, n_bdim] = unwrapTensorAtLevel(n, cur_level);
  25837. auto results = batch_rule(x_value, x_bdim, n_value, n_bdim);
  25838. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25839. }
  25840. template <typename batch_rule_t, batch_rule_t batch_rule>
  25841. at::Tensor special_shifted_chebyshev_polynomial_v_x_scalar_generated_plumbing(const at::Scalar & x, const at::Tensor & n) {
  25842. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25843. auto maybe_layer = maybeCurrentDynamicLayer();
  25844. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25845. int64_t cur_level = maybe_layer->layerId();
  25846. if (!isBatchedAtLevel(n, cur_level)) {
  25847. return at::_ops::special_shifted_chebyshev_polynomial_v_x_scalar::call(x, n);
  25848. }
  25849. auto [n_value, n_bdim] = unwrapTensorAtLevel(n, cur_level);
  25850. auto results = batch_rule(x, n_value, n_bdim);
  25851. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25852. }
  25853. template <typename batch_rule_t, batch_rule_t batch_rule>
  25854. at::Tensor special_shifted_chebyshev_polynomial_v_n_scalar_generated_plumbing(const at::Tensor & x, const at::Scalar & n) {
  25855. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25856. auto maybe_layer = maybeCurrentDynamicLayer();
  25857. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25858. int64_t cur_level = maybe_layer->layerId();
  25859. if (!isBatchedAtLevel(x, cur_level)) {
  25860. return at::_ops::special_shifted_chebyshev_polynomial_v_n_scalar::call(x, n);
  25861. }
  25862. auto [x_value, x_bdim] = unwrapTensorAtLevel(x, cur_level);
  25863. auto results = batch_rule(x_value, x_bdim, n);
  25864. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25865. }
  25866. template <typename batch_rule_t, batch_rule_t batch_rule>
  25867. at::Tensor special_shifted_chebyshev_polynomial_w_generated_plumbing(const at::Tensor & x, const at::Tensor & n) {
  25868. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25869. auto maybe_layer = maybeCurrentDynamicLayer();
  25870. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25871. int64_t cur_level = maybe_layer->layerId();
  25872. if (!isBatchedAtLevel(x, cur_level) && !isBatchedAtLevel(n, cur_level)) {
  25873. return at::_ops::special_shifted_chebyshev_polynomial_w::call(x, n);
  25874. }
  25875. auto [x_value, x_bdim] = unwrapTensorAtLevel(x, cur_level);
  25876. auto [n_value, n_bdim] = unwrapTensorAtLevel(n, cur_level);
  25877. auto results = batch_rule(x_value, x_bdim, n_value, n_bdim);
  25878. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25879. }
  25880. template <typename batch_rule_t, batch_rule_t batch_rule>
  25881. at::Tensor special_shifted_chebyshev_polynomial_w_x_scalar_generated_plumbing(const at::Scalar & x, const at::Tensor & n) {
  25882. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25883. auto maybe_layer = maybeCurrentDynamicLayer();
  25884. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25885. int64_t cur_level = maybe_layer->layerId();
  25886. if (!isBatchedAtLevel(n, cur_level)) {
  25887. return at::_ops::special_shifted_chebyshev_polynomial_w_x_scalar::call(x, n);
  25888. }
  25889. auto [n_value, n_bdim] = unwrapTensorAtLevel(n, cur_level);
  25890. auto results = batch_rule(x, n_value, n_bdim);
  25891. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25892. }
  25893. template <typename batch_rule_t, batch_rule_t batch_rule>
  25894. at::Tensor special_shifted_chebyshev_polynomial_w_n_scalar_generated_plumbing(const at::Tensor & x, const at::Scalar & n) {
  25895. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25896. auto maybe_layer = maybeCurrentDynamicLayer();
  25897. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25898. int64_t cur_level = maybe_layer->layerId();
  25899. if (!isBatchedAtLevel(x, cur_level)) {
  25900. return at::_ops::special_shifted_chebyshev_polynomial_w_n_scalar::call(x, n);
  25901. }
  25902. auto [x_value, x_bdim] = unwrapTensorAtLevel(x, cur_level);
  25903. auto results = batch_rule(x_value, x_bdim, n);
  25904. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25905. }
  25906. template <typename batch_rule_t, batch_rule_t batch_rule>
  25907. at::Tensor special_spherical_bessel_j0_generated_plumbing(const at::Tensor & x) {
  25908. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25909. auto maybe_layer = maybeCurrentDynamicLayer();
  25910. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25911. int64_t cur_level = maybe_layer->layerId();
  25912. if (!isBatchedAtLevel(x, cur_level)) {
  25913. return at::_ops::special_spherical_bessel_j0::call(x);
  25914. }
  25915. auto [x_value, x_bdim] = unwrapTensorAtLevel(x, cur_level);
  25916. auto results = batch_rule(x_value, x_bdim);
  25917. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25918. }
  25919. template <typename batch_rule_t, batch_rule_t batch_rule>
  25920. at::Tensor _foobar_generated_plumbing(const at::Tensor & self, bool arg1, bool arg2, bool arg3) {
  25921. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25922. auto maybe_layer = maybeCurrentDynamicLayer();
  25923. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  25924. int64_t cur_level = maybe_layer->layerId();
  25925. if (!isBatchedAtLevel(self, cur_level)) {
  25926. return at::_ops::_foobar::call(self, arg1, arg2, arg3);
  25927. }
  25928. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  25929. auto results = batch_rule(self_value, self_bdim, arg1, arg2, arg3);
  25930. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  25931. }
  25932. template <typename batch_rule_t, batch_rule_t batch_rule>
  25933. void _fused_adam__generated_plumbing(at::TensorList self, at::TensorList grads, at::TensorList exp_avgs, at::TensorList exp_avg_sqs, at::TensorList max_exp_avg_sqs, at::TensorList state_steps, double lr, double beta1, double beta2, double weight_decay, double eps, bool amsgrad, bool maximize, const ::std::optional<at::Tensor> & grad_scale, const ::std::optional<at::Tensor> & found_inf) {
  25934. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25935. auto maybe_layer = maybeCurrentDynamicLayer();
  25936. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  25937. int64_t cur_level = maybe_layer->layerId();
  25938. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(grads, cur_level) && !isBatchedAtLevel(exp_avgs, cur_level) && !isBatchedAtLevel(exp_avg_sqs, cur_level) && !isBatchedAtLevel(max_exp_avg_sqs, cur_level) && !isBatchedAtLevel(state_steps, cur_level) && !isBatchedAtLevel(grad_scale, cur_level) && !isBatchedAtLevel(found_inf, cur_level)) {
  25939. return at::_ops::_fused_adam_::call(self, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps, lr, beta1, beta2, weight_decay, eps, amsgrad, maximize, grad_scale, found_inf);
  25940. }
  25941. std::optional<Tensor> grad_scale_value;
  25942. std::optional<int64_t> grad_scale_bdim;
  25943. if (grad_scale) {
  25944. std::tie(grad_scale_value, grad_scale_bdim) = unwrapTensorAtLevel(grad_scale.value(), cur_level);
  25945. }
  25946. std::optional<Tensor> found_inf_value;
  25947. std::optional<int64_t> found_inf_bdim;
  25948. if (found_inf) {
  25949. std::tie(found_inf_value, found_inf_bdim) = unwrapTensorAtLevel(found_inf.value(), cur_level);
  25950. }
  25951. batch_rule(self, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps, lr, beta1, beta2, weight_decay, eps, amsgrad, maximize, grad_scale_value, grad_scale_bdim, found_inf_value, found_inf_bdim);
  25952. }
  25953. template <typename batch_rule_t, batch_rule_t batch_rule>
  25954. void _fused_adam__tensor_lr_generated_plumbing(at::TensorList self, at::TensorList grads, at::TensorList exp_avgs, at::TensorList exp_avg_sqs, at::TensorList max_exp_avg_sqs, at::TensorList state_steps, const at::Tensor & lr, double beta1, double beta2, double weight_decay, double eps, bool amsgrad, bool maximize, const ::std::optional<at::Tensor> & grad_scale, const ::std::optional<at::Tensor> & found_inf) {
  25955. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25956. auto maybe_layer = maybeCurrentDynamicLayer();
  25957. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  25958. int64_t cur_level = maybe_layer->layerId();
  25959. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(grads, cur_level) && !isBatchedAtLevel(exp_avgs, cur_level) && !isBatchedAtLevel(exp_avg_sqs, cur_level) && !isBatchedAtLevel(max_exp_avg_sqs, cur_level) && !isBatchedAtLevel(state_steps, cur_level) && !isBatchedAtLevel(lr, cur_level) && !isBatchedAtLevel(grad_scale, cur_level) && !isBatchedAtLevel(found_inf, cur_level)) {
  25960. return at::_ops::_fused_adam__tensor_lr::call(self, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps, lr, beta1, beta2, weight_decay, eps, amsgrad, maximize, grad_scale, found_inf);
  25961. }
  25962. auto [lr_value, lr_bdim] = unwrapTensorAtLevel(lr, cur_level);
  25963. std::optional<Tensor> grad_scale_value;
  25964. std::optional<int64_t> grad_scale_bdim;
  25965. if (grad_scale) {
  25966. std::tie(grad_scale_value, grad_scale_bdim) = unwrapTensorAtLevel(grad_scale.value(), cur_level);
  25967. }
  25968. std::optional<Tensor> found_inf_value;
  25969. std::optional<int64_t> found_inf_bdim;
  25970. if (found_inf) {
  25971. std::tie(found_inf_value, found_inf_bdim) = unwrapTensorAtLevel(found_inf.value(), cur_level);
  25972. }
  25973. batch_rule(self, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps, lr_value, lr_bdim, beta1, beta2, weight_decay, eps, amsgrad, maximize, grad_scale_value, grad_scale_bdim, found_inf_value, found_inf_bdim);
  25974. }
  25975. template <typename batch_rule_t, batch_rule_t batch_rule>
  25976. void _fused_adamw__generated_plumbing(at::TensorList self, at::TensorList grads, at::TensorList exp_avgs, at::TensorList exp_avg_sqs, at::TensorList max_exp_avg_sqs, at::TensorList state_steps, double lr, double beta1, double beta2, double weight_decay, double eps, bool amsgrad, bool maximize, const ::std::optional<at::Tensor> & grad_scale, const ::std::optional<at::Tensor> & found_inf) {
  25977. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25978. auto maybe_layer = maybeCurrentDynamicLayer();
  25979. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  25980. int64_t cur_level = maybe_layer->layerId();
  25981. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(grads, cur_level) && !isBatchedAtLevel(exp_avgs, cur_level) && !isBatchedAtLevel(exp_avg_sqs, cur_level) && !isBatchedAtLevel(max_exp_avg_sqs, cur_level) && !isBatchedAtLevel(state_steps, cur_level) && !isBatchedAtLevel(grad_scale, cur_level) && !isBatchedAtLevel(found_inf, cur_level)) {
  25982. return at::_ops::_fused_adamw_::call(self, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps, lr, beta1, beta2, weight_decay, eps, amsgrad, maximize, grad_scale, found_inf);
  25983. }
  25984. std::optional<Tensor> grad_scale_value;
  25985. std::optional<int64_t> grad_scale_bdim;
  25986. if (grad_scale) {
  25987. std::tie(grad_scale_value, grad_scale_bdim) = unwrapTensorAtLevel(grad_scale.value(), cur_level);
  25988. }
  25989. std::optional<Tensor> found_inf_value;
  25990. std::optional<int64_t> found_inf_bdim;
  25991. if (found_inf) {
  25992. std::tie(found_inf_value, found_inf_bdim) = unwrapTensorAtLevel(found_inf.value(), cur_level);
  25993. }
  25994. batch_rule(self, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps, lr, beta1, beta2, weight_decay, eps, amsgrad, maximize, grad_scale_value, grad_scale_bdim, found_inf_value, found_inf_bdim);
  25995. }
  25996. template <typename batch_rule_t, batch_rule_t batch_rule>
  25997. void _fused_adamw__tensor_lr_generated_plumbing(at::TensorList self, at::TensorList grads, at::TensorList exp_avgs, at::TensorList exp_avg_sqs, at::TensorList max_exp_avg_sqs, at::TensorList state_steps, const at::Tensor & lr, double beta1, double beta2, double weight_decay, double eps, bool amsgrad, bool maximize, const ::std::optional<at::Tensor> & grad_scale, const ::std::optional<at::Tensor> & found_inf) {
  25998. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  25999. auto maybe_layer = maybeCurrentDynamicLayer();
  26000. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  26001. int64_t cur_level = maybe_layer->layerId();
  26002. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(grads, cur_level) && !isBatchedAtLevel(exp_avgs, cur_level) && !isBatchedAtLevel(exp_avg_sqs, cur_level) && !isBatchedAtLevel(max_exp_avg_sqs, cur_level) && !isBatchedAtLevel(state_steps, cur_level) && !isBatchedAtLevel(lr, cur_level) && !isBatchedAtLevel(grad_scale, cur_level) && !isBatchedAtLevel(found_inf, cur_level)) {
  26003. return at::_ops::_fused_adamw__tensor_lr::call(self, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps, lr, beta1, beta2, weight_decay, eps, amsgrad, maximize, grad_scale, found_inf);
  26004. }
  26005. auto [lr_value, lr_bdim] = unwrapTensorAtLevel(lr, cur_level);
  26006. std::optional<Tensor> grad_scale_value;
  26007. std::optional<int64_t> grad_scale_bdim;
  26008. if (grad_scale) {
  26009. std::tie(grad_scale_value, grad_scale_bdim) = unwrapTensorAtLevel(grad_scale.value(), cur_level);
  26010. }
  26011. std::optional<Tensor> found_inf_value;
  26012. std::optional<int64_t> found_inf_bdim;
  26013. if (found_inf) {
  26014. std::tie(found_inf_value, found_inf_bdim) = unwrapTensorAtLevel(found_inf.value(), cur_level);
  26015. }
  26016. batch_rule(self, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps, lr_value, lr_bdim, beta1, beta2, weight_decay, eps, amsgrad, maximize, grad_scale_value, grad_scale_bdim, found_inf_value, found_inf_bdim);
  26017. }
  26018. template <typename batch_rule_t, batch_rule_t batch_rule>
  26019. void _fused_sgd__generated_plumbing(at::TensorList self, at::TensorList grads, at::TensorList momentum_buffer_list, double weight_decay, double momentum, double lr, double dampening, bool nesterov, bool maximize, bool is_first_step, const ::std::optional<at::Tensor> & grad_scale, const ::std::optional<at::Tensor> & found_inf) {
  26020. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26021. auto maybe_layer = maybeCurrentDynamicLayer();
  26022. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  26023. int64_t cur_level = maybe_layer->layerId();
  26024. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(grads, cur_level) && !isBatchedAtLevel(momentum_buffer_list, cur_level) && !isBatchedAtLevel(grad_scale, cur_level) && !isBatchedAtLevel(found_inf, cur_level)) {
  26025. return at::_ops::_fused_sgd_::call(self, grads, momentum_buffer_list, weight_decay, momentum, lr, dampening, nesterov, maximize, is_first_step, grad_scale, found_inf);
  26026. }
  26027. std::optional<Tensor> grad_scale_value;
  26028. std::optional<int64_t> grad_scale_bdim;
  26029. if (grad_scale) {
  26030. std::tie(grad_scale_value, grad_scale_bdim) = unwrapTensorAtLevel(grad_scale.value(), cur_level);
  26031. }
  26032. std::optional<Tensor> found_inf_value;
  26033. std::optional<int64_t> found_inf_bdim;
  26034. if (found_inf) {
  26035. std::tie(found_inf_value, found_inf_bdim) = unwrapTensorAtLevel(found_inf.value(), cur_level);
  26036. }
  26037. batch_rule(self, grads, momentum_buffer_list, weight_decay, momentum, lr, dampening, nesterov, maximize, is_first_step, grad_scale_value, grad_scale_bdim, found_inf_value, found_inf_bdim);
  26038. }
  26039. template <typename batch_rule_t, batch_rule_t batch_rule>
  26040. void _fused_sgd__tensor_lr_generated_plumbing(at::TensorList self, at::TensorList grads, at::TensorList momentum_buffer_list, double weight_decay, double momentum, const at::Tensor & lr, double dampening, bool nesterov, bool maximize, bool is_first_step, const ::std::optional<at::Tensor> & grad_scale, const ::std::optional<at::Tensor> & found_inf) {
  26041. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26042. auto maybe_layer = maybeCurrentDynamicLayer();
  26043. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  26044. int64_t cur_level = maybe_layer->layerId();
  26045. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(grads, cur_level) && !isBatchedAtLevel(momentum_buffer_list, cur_level) && !isBatchedAtLevel(lr, cur_level) && !isBatchedAtLevel(grad_scale, cur_level) && !isBatchedAtLevel(found_inf, cur_level)) {
  26046. return at::_ops::_fused_sgd__tensor_lr::call(self, grads, momentum_buffer_list, weight_decay, momentum, lr, dampening, nesterov, maximize, is_first_step, grad_scale, found_inf);
  26047. }
  26048. auto [lr_value, lr_bdim] = unwrapTensorAtLevel(lr, cur_level);
  26049. std::optional<Tensor> grad_scale_value;
  26050. std::optional<int64_t> grad_scale_bdim;
  26051. if (grad_scale) {
  26052. std::tie(grad_scale_value, grad_scale_bdim) = unwrapTensorAtLevel(grad_scale.value(), cur_level);
  26053. }
  26054. std::optional<Tensor> found_inf_value;
  26055. std::optional<int64_t> found_inf_bdim;
  26056. if (found_inf) {
  26057. std::tie(found_inf_value, found_inf_bdim) = unwrapTensorAtLevel(found_inf.value(), cur_level);
  26058. }
  26059. batch_rule(self, grads, momentum_buffer_list, weight_decay, momentum, lr_value, lr_bdim, dampening, nesterov, maximize, is_first_step, grad_scale_value, grad_scale_bdim, found_inf_value, found_inf_bdim);
  26060. }
  26061. template <typename batch_rule_t, batch_rule_t batch_rule>
  26062. void _fused_adagrad__generated_plumbing(at::TensorList self, at::TensorList grads, at::TensorList state_sums, at::TensorList state_steps, double lr, double lr_decay, double weight_decay, double eps, bool maximize, const ::std::optional<at::Tensor> & grad_scale, const ::std::optional<at::Tensor> & found_inf) {
  26063. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26064. auto maybe_layer = maybeCurrentDynamicLayer();
  26065. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  26066. int64_t cur_level = maybe_layer->layerId();
  26067. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(grads, cur_level) && !isBatchedAtLevel(state_sums, cur_level) && !isBatchedAtLevel(state_steps, cur_level) && !isBatchedAtLevel(grad_scale, cur_level) && !isBatchedAtLevel(found_inf, cur_level)) {
  26068. return at::_ops::_fused_adagrad_::call(self, grads, state_sums, state_steps, lr, lr_decay, weight_decay, eps, maximize, grad_scale, found_inf);
  26069. }
  26070. std::optional<Tensor> grad_scale_value;
  26071. std::optional<int64_t> grad_scale_bdim;
  26072. if (grad_scale) {
  26073. std::tie(grad_scale_value, grad_scale_bdim) = unwrapTensorAtLevel(grad_scale.value(), cur_level);
  26074. }
  26075. std::optional<Tensor> found_inf_value;
  26076. std::optional<int64_t> found_inf_bdim;
  26077. if (found_inf) {
  26078. std::tie(found_inf_value, found_inf_bdim) = unwrapTensorAtLevel(found_inf.value(), cur_level);
  26079. }
  26080. batch_rule(self, grads, state_sums, state_steps, lr, lr_decay, weight_decay, eps, maximize, grad_scale_value, grad_scale_bdim, found_inf_value, found_inf_bdim);
  26081. }
  26082. template <typename batch_rule_t, batch_rule_t batch_rule>
  26083. void _fused_adagrad__tensor_lr_generated_plumbing(at::TensorList self, at::TensorList grads, at::TensorList state_sums, at::TensorList state_steps, const at::Tensor & lr, double lr_decay, double weight_decay, double eps, bool maximize, const ::std::optional<at::Tensor> & grad_scale, const ::std::optional<at::Tensor> & found_inf) {
  26084. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26085. auto maybe_layer = maybeCurrentDynamicLayer();
  26086. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  26087. int64_t cur_level = maybe_layer->layerId();
  26088. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(grads, cur_level) && !isBatchedAtLevel(state_sums, cur_level) && !isBatchedAtLevel(state_steps, cur_level) && !isBatchedAtLevel(lr, cur_level) && !isBatchedAtLevel(grad_scale, cur_level) && !isBatchedAtLevel(found_inf, cur_level)) {
  26089. return at::_ops::_fused_adagrad__tensor_lr::call(self, grads, state_sums, state_steps, lr, lr_decay, weight_decay, eps, maximize, grad_scale, found_inf);
  26090. }
  26091. auto [lr_value, lr_bdim] = unwrapTensorAtLevel(lr, cur_level);
  26092. std::optional<Tensor> grad_scale_value;
  26093. std::optional<int64_t> grad_scale_bdim;
  26094. if (grad_scale) {
  26095. std::tie(grad_scale_value, grad_scale_bdim) = unwrapTensorAtLevel(grad_scale.value(), cur_level);
  26096. }
  26097. std::optional<Tensor> found_inf_value;
  26098. std::optional<int64_t> found_inf_bdim;
  26099. if (found_inf) {
  26100. std::tie(found_inf_value, found_inf_bdim) = unwrapTensorAtLevel(found_inf.value(), cur_level);
  26101. }
  26102. batch_rule(self, grads, state_sums, state_steps, lr_value, lr_bdim, lr_decay, weight_decay, eps, maximize, grad_scale_value, grad_scale_bdim, found_inf_value, found_inf_bdim);
  26103. }
  26104. template <typename batch_rule_t, batch_rule_t batch_rule>
  26105. void _propagate_xla_data_generated_plumbing(const at::Tensor & input, const at::Tensor & output) {
  26106. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26107. auto maybe_layer = maybeCurrentDynamicLayer();
  26108. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  26109. int64_t cur_level = maybe_layer->layerId();
  26110. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(output, cur_level)) {
  26111. return at::_ops::_propagate_xla_data::call(input, output);
  26112. }
  26113. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  26114. auto [output_value, output_bdim] = unwrapTensorAtLevel(output, cur_level);
  26115. batch_rule(input_value, input_bdim, output_value, output_bdim);
  26116. }
  26117. template <typename batch_rule_t, batch_rule_t batch_rule>
  26118. void _cudnn_rnn_backward_out_generated_plumbing(const at::Tensor & input, at::TensorList weight, int64_t weight_stride0, const at::Tensor & weight_buf, const at::Tensor & hx, const ::std::optional<at::Tensor> & cx, const at::Tensor & output, const ::std::optional<at::Tensor> & grad_output, const ::std::optional<at::Tensor> & grad_hy, const ::std::optional<at::Tensor> & grad_cy, int64_t mode, c10::SymInt hidden_size, c10::SymInt proj_size, int64_t num_layers, bool batch_first, double dropout, bool train, bool bidirectional, c10::SymIntArrayRef batch_sizes, const ::std::optional<at::Tensor> & dropout_state, const at::Tensor & reserve, ::std::array<bool,4> output_mask, at::Tensor & out0, at::Tensor & out1, at::Tensor & out2, at::TensorList out3) {
  26119. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26120. auto maybe_layer = maybeCurrentDynamicLayer();
  26121. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  26122. int64_t cur_level = maybe_layer->layerId();
  26123. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(weight_buf, cur_level) && !isBatchedAtLevel(hx, cur_level) && !isBatchedAtLevel(cx, cur_level) && !isBatchedAtLevel(output, cur_level) && !isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(grad_hy, cur_level) && !isBatchedAtLevel(grad_cy, cur_level) && !isBatchedAtLevel(dropout_state, cur_level) && !isBatchedAtLevel(reserve, cur_level) && !isBatchedAtLevel(out0, cur_level) && !isBatchedAtLevel(out1, cur_level) && !isBatchedAtLevel(out2, cur_level) && !isBatchedAtLevel(out3, cur_level)) {
  26124. return at::_ops::_cudnn_rnn_backward_out::call(input, weight, weight_stride0, weight_buf, hx, cx, output, grad_output, grad_hy, grad_cy, mode, hidden_size, proj_size, num_layers, batch_first, dropout, train, bidirectional, batch_sizes, dropout_state, reserve, output_mask, out0, out1, out2, out3);
  26125. }
  26126. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  26127. auto [weight_buf_value, weight_buf_bdim] = unwrapTensorAtLevel(weight_buf, cur_level);
  26128. auto [hx_value, hx_bdim] = unwrapTensorAtLevel(hx, cur_level);
  26129. auto [output_value, output_bdim] = unwrapTensorAtLevel(output, cur_level);
  26130. auto [reserve_value, reserve_bdim] = unwrapTensorAtLevel(reserve, cur_level);
  26131. auto [out0_value, out0_bdim] = unwrapTensorAtLevel(out0, cur_level);
  26132. auto [out1_value, out1_bdim] = unwrapTensorAtLevel(out1, cur_level);
  26133. auto [out2_value, out2_bdim] = unwrapTensorAtLevel(out2, cur_level);
  26134. std::optional<Tensor> cx_value;
  26135. std::optional<int64_t> cx_bdim;
  26136. if (cx) {
  26137. std::tie(cx_value, cx_bdim) = unwrapTensorAtLevel(cx.value(), cur_level);
  26138. }
  26139. std::optional<Tensor> grad_output_value;
  26140. std::optional<int64_t> grad_output_bdim;
  26141. if (grad_output) {
  26142. std::tie(grad_output_value, grad_output_bdim) = unwrapTensorAtLevel(grad_output.value(), cur_level);
  26143. }
  26144. std::optional<Tensor> grad_hy_value;
  26145. std::optional<int64_t> grad_hy_bdim;
  26146. if (grad_hy) {
  26147. std::tie(grad_hy_value, grad_hy_bdim) = unwrapTensorAtLevel(grad_hy.value(), cur_level);
  26148. }
  26149. std::optional<Tensor> grad_cy_value;
  26150. std::optional<int64_t> grad_cy_bdim;
  26151. if (grad_cy) {
  26152. std::tie(grad_cy_value, grad_cy_bdim) = unwrapTensorAtLevel(grad_cy.value(), cur_level);
  26153. }
  26154. std::optional<Tensor> dropout_state_value;
  26155. std::optional<int64_t> dropout_state_bdim;
  26156. if (dropout_state) {
  26157. std::tie(dropout_state_value, dropout_state_bdim) = unwrapTensorAtLevel(dropout_state.value(), cur_level);
  26158. }
  26159. batch_rule(input_value, input_bdim, weight, weight_stride0, weight_buf_value, weight_buf_bdim, hx_value, hx_bdim, cx_value, cx_bdim, output_value, output_bdim, grad_output_value, grad_output_bdim, grad_hy_value, grad_hy_bdim, grad_cy_value, grad_cy_bdim, mode, hidden_size, proj_size, num_layers, batch_first, dropout, train, bidirectional, batch_sizes, dropout_state_value, dropout_state_bdim, reserve_value, reserve_bdim, output_mask, out0_value, out0_bdim, out1_value, out1_bdim, out2_value, out2_bdim, out3);
  26160. }
  26161. template <typename batch_rule_t, batch_rule_t batch_rule>
  26162. at::Tensor bernoulli_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & p, ::std::optional<at::Generator> generator) {
  26163. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26164. auto maybe_layer = maybeCurrentDynamicLayer();
  26165. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  26166. int64_t cur_level = maybe_layer->layerId();
  26167. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(p, cur_level)) {
  26168. return at::_ops::bernoulli_Tensor::call(self, p, generator);
  26169. }
  26170. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  26171. auto [p_value, p_bdim] = unwrapTensorAtLevel(p, cur_level);
  26172. auto results = batch_rule(self_value, self_bdim, p_value, p_bdim, generator);
  26173. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  26174. }
  26175. template <typename batch_rule_t, batch_rule_t batch_rule>
  26176. at::Tensor embedding_renorm_generated_plumbing(const at::Tensor & self, const at::Tensor & indices, double max_norm, double norm_type) {
  26177. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26178. auto maybe_layer = maybeCurrentDynamicLayer();
  26179. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  26180. int64_t cur_level = maybe_layer->layerId();
  26181. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(indices, cur_level)) {
  26182. return at::_ops::embedding_renorm::call(self, indices, max_norm, norm_type);
  26183. }
  26184. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  26185. auto [indices_value, indices_bdim] = unwrapTensorAtLevel(indices, cur_level);
  26186. auto results = batch_rule(self_value, self_bdim, indices_value, indices_bdim, max_norm, norm_type);
  26187. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  26188. }
  26189. template <typename batch_rule_t, batch_rule_t batch_rule>
  26190. at::Tensor resize_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef size, ::std::optional<at::MemoryFormat> memory_format) {
  26191. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26192. auto maybe_layer = maybeCurrentDynamicLayer();
  26193. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  26194. int64_t cur_level = maybe_layer->layerId();
  26195. if (!isBatchedAtLevel(self, cur_level)) {
  26196. return at::_ops::resize::call(self, size, memory_format);
  26197. }
  26198. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  26199. auto results = batch_rule(self_value, self_bdim, size, memory_format);
  26200. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  26201. }
  26202. template <typename batch_rule_t, batch_rule_t batch_rule>
  26203. at::Tensor _resize_output_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef size, at::Device device) {
  26204. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26205. auto maybe_layer = maybeCurrentDynamicLayer();
  26206. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  26207. int64_t cur_level = maybe_layer->layerId();
  26208. if (!isBatchedAtLevel(self, cur_level)) {
  26209. return at::_ops::_resize_output::call(self, size, device);
  26210. }
  26211. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  26212. auto results = batch_rule(self_value, self_bdim, size, device);
  26213. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  26214. }
  26215. template <typename batch_rule_t, batch_rule_t batch_rule>
  26216. at::Tensor _index_put_impl_generated_plumbing(const at::Tensor & self, const c10::List<::std::optional<at::Tensor>> & indices, const at::Tensor & values, bool accumulate, bool unsafe) {
  26217. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26218. auto maybe_layer = maybeCurrentDynamicLayer();
  26219. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  26220. int64_t cur_level = maybe_layer->layerId();
  26221. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(indices, cur_level) && !isBatchedAtLevel(values, cur_level)) {
  26222. return at::_ops::_index_put_impl::call(self, indices, values, accumulate, unsafe);
  26223. }
  26224. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  26225. auto [values_value, values_bdim] = unwrapTensorAtLevel(values, cur_level);
  26226. auto results = batch_rule(self_value, self_bdim, indices, values_value, values_bdim, accumulate, unsafe);
  26227. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  26228. }
  26229. template <typename batch_rule_t, batch_rule_t batch_rule>
  26230. void miopen_rnn_backward_out_generated_plumbing(const at::Tensor & input, at::TensorList weight, int64_t weight_stride0, const at::Tensor & weight_buf, const at::Tensor & hx, const ::std::optional<at::Tensor> & cx, const at::Tensor & output, const ::std::optional<at::Tensor> & grad_output, const ::std::optional<at::Tensor> & grad_hy, const ::std::optional<at::Tensor> & grad_cy, int64_t mode, int64_t hidden_size, int64_t num_layers, bool batch_first, double dropout, bool train, bool bidirectional, at::IntArrayRef batch_sizes, const ::std::optional<at::Tensor> & dropout_state, const at::Tensor & reserve, ::std::array<bool,4> output_mask, at::Tensor & out0, at::Tensor & out1, at::Tensor & out2, at::TensorList out3) {
  26231. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26232. auto maybe_layer = maybeCurrentDynamicLayer();
  26233. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  26234. int64_t cur_level = maybe_layer->layerId();
  26235. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(weight_buf, cur_level) && !isBatchedAtLevel(hx, cur_level) && !isBatchedAtLevel(cx, cur_level) && !isBatchedAtLevel(output, cur_level) && !isBatchedAtLevel(grad_output, cur_level) && !isBatchedAtLevel(grad_hy, cur_level) && !isBatchedAtLevel(grad_cy, cur_level) && !isBatchedAtLevel(dropout_state, cur_level) && !isBatchedAtLevel(reserve, cur_level) && !isBatchedAtLevel(out0, cur_level) && !isBatchedAtLevel(out1, cur_level) && !isBatchedAtLevel(out2, cur_level) && !isBatchedAtLevel(out3, cur_level)) {
  26236. return at::_ops::miopen_rnn_backward_out::call(input, weight, weight_stride0, weight_buf, hx, cx, output, grad_output, grad_hy, grad_cy, mode, hidden_size, num_layers, batch_first, dropout, train, bidirectional, batch_sizes, dropout_state, reserve, output_mask, out0, out1, out2, out3);
  26237. }
  26238. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  26239. auto [weight_buf_value, weight_buf_bdim] = unwrapTensorAtLevel(weight_buf, cur_level);
  26240. auto [hx_value, hx_bdim] = unwrapTensorAtLevel(hx, cur_level);
  26241. auto [output_value, output_bdim] = unwrapTensorAtLevel(output, cur_level);
  26242. auto [reserve_value, reserve_bdim] = unwrapTensorAtLevel(reserve, cur_level);
  26243. auto [out0_value, out0_bdim] = unwrapTensorAtLevel(out0, cur_level);
  26244. auto [out1_value, out1_bdim] = unwrapTensorAtLevel(out1, cur_level);
  26245. auto [out2_value, out2_bdim] = unwrapTensorAtLevel(out2, cur_level);
  26246. std::optional<Tensor> cx_value;
  26247. std::optional<int64_t> cx_bdim;
  26248. if (cx) {
  26249. std::tie(cx_value, cx_bdim) = unwrapTensorAtLevel(cx.value(), cur_level);
  26250. }
  26251. std::optional<Tensor> grad_output_value;
  26252. std::optional<int64_t> grad_output_bdim;
  26253. if (grad_output) {
  26254. std::tie(grad_output_value, grad_output_bdim) = unwrapTensorAtLevel(grad_output.value(), cur_level);
  26255. }
  26256. std::optional<Tensor> grad_hy_value;
  26257. std::optional<int64_t> grad_hy_bdim;
  26258. if (grad_hy) {
  26259. std::tie(grad_hy_value, grad_hy_bdim) = unwrapTensorAtLevel(grad_hy.value(), cur_level);
  26260. }
  26261. std::optional<Tensor> grad_cy_value;
  26262. std::optional<int64_t> grad_cy_bdim;
  26263. if (grad_cy) {
  26264. std::tie(grad_cy_value, grad_cy_bdim) = unwrapTensorAtLevel(grad_cy.value(), cur_level);
  26265. }
  26266. std::optional<Tensor> dropout_state_value;
  26267. std::optional<int64_t> dropout_state_bdim;
  26268. if (dropout_state) {
  26269. std::tie(dropout_state_value, dropout_state_bdim) = unwrapTensorAtLevel(dropout_state.value(), cur_level);
  26270. }
  26271. batch_rule(input_value, input_bdim, weight, weight_stride0, weight_buf_value, weight_buf_bdim, hx_value, hx_bdim, cx_value, cx_bdim, output_value, output_bdim, grad_output_value, grad_output_bdim, grad_hy_value, grad_hy_bdim, grad_cy_value, grad_cy_bdim, mode, hidden_size, num_layers, batch_first, dropout, train, bidirectional, batch_sizes, dropout_state_value, dropout_state_bdim, reserve_value, reserve_bdim, output_mask, out0_value, out0_bdim, out1_value, out1_bdim, out2_value, out2_bdim, out3);
  26272. }
  26273. template <typename batch_rule_t, batch_rule_t batch_rule>
  26274. ::std::tuple<at::Tensor,at::Tensor,at::Tensor,at::Tensor,at::Tensor> _native_batch_norm_legit_functional_generated_plumbing(const at::Tensor & input, const ::std::optional<at::Tensor> & weight, const ::std::optional<at::Tensor> & bias, const at::Tensor & running_mean, const at::Tensor & running_var, bool training, double momentum, double eps) {
  26275. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26276. auto maybe_layer = maybeCurrentDynamicLayer();
  26277. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  26278. int64_t cur_level = maybe_layer->layerId();
  26279. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level) && !isBatchedAtLevel(running_mean, cur_level) && !isBatchedAtLevel(running_var, cur_level)) {
  26280. return at::_ops::_native_batch_norm_legit_functional::call(input, weight, bias, running_mean, running_var, training, momentum, eps);
  26281. }
  26282. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  26283. auto [running_mean_value, running_mean_bdim] = unwrapTensorAtLevel(running_mean, cur_level);
  26284. auto [running_var_value, running_var_bdim] = unwrapTensorAtLevel(running_var, cur_level);
  26285. std::optional<Tensor> weight_value;
  26286. std::optional<int64_t> weight_bdim;
  26287. if (weight) {
  26288. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  26289. }
  26290. std::optional<Tensor> bias_value;
  26291. std::optional<int64_t> bias_bdim;
  26292. if (bias) {
  26293. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  26294. }
  26295. auto results = batch_rule(input_value, input_bdim, weight_value, weight_bdim, bias_value, bias_bdim, running_mean_value, running_mean_bdim, running_var_value, running_var_bdim, training, momentum, eps);
  26296. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level), makeBatched(std::get<6>(results), std::get<7>(results), cur_level), makeBatched(std::get<8>(results), std::get<9>(results), cur_level));
  26297. }
  26298. template <typename batch_rule_t, batch_rule_t batch_rule>
  26299. void unsafe_split_Tensor_out_generated_plumbing(const at::Tensor & self, c10::SymInt split_size, int64_t dim, at::TensorList out) {
  26300. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26301. auto maybe_layer = maybeCurrentDynamicLayer();
  26302. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  26303. int64_t cur_level = maybe_layer->layerId();
  26304. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  26305. return at::_ops::unsafe_split_Tensor_out::call(self, split_size, dim, out);
  26306. }
  26307. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  26308. batch_rule(self_value, self_bdim, split_size, dim, out);
  26309. }
  26310. template <typename batch_rule_t, batch_rule_t batch_rule>
  26311. void unsafe_split_with_sizes_out_generated_plumbing(const at::Tensor & self, c10::SymIntArrayRef split_sizes, int64_t dim, at::TensorList out) {
  26312. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26313. auto maybe_layer = maybeCurrentDynamicLayer();
  26314. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  26315. int64_t cur_level = maybe_layer->layerId();
  26316. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  26317. return at::_ops::unsafe_split_with_sizes_out::call(self, split_sizes, dim, out);
  26318. }
  26319. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  26320. batch_rule(self_value, self_bdim, split_sizes, dim, out);
  26321. }
  26322. template <typename batch_rule_t, batch_rule_t batch_rule>
  26323. ::std::tuple<at::Tensor,at::Tensor,at::Tensor,at::Tensor,at::Tensor,at::Tensor> _batch_norm_with_update_functional_generated_plumbing(const at::Tensor & input, const ::std::optional<at::Tensor> & weight, const ::std::optional<at::Tensor> & bias, const at::Tensor & running_mean, const at::Tensor & running_var, double momentum, double eps) {
  26324. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26325. auto maybe_layer = maybeCurrentDynamicLayer();
  26326. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  26327. int64_t cur_level = maybe_layer->layerId();
  26328. if (!isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(bias, cur_level) && !isBatchedAtLevel(running_mean, cur_level) && !isBatchedAtLevel(running_var, cur_level)) {
  26329. return at::_ops::_batch_norm_with_update_functional::call(input, weight, bias, running_mean, running_var, momentum, eps);
  26330. }
  26331. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  26332. auto [running_mean_value, running_mean_bdim] = unwrapTensorAtLevel(running_mean, cur_level);
  26333. auto [running_var_value, running_var_bdim] = unwrapTensorAtLevel(running_var, cur_level);
  26334. std::optional<Tensor> weight_value;
  26335. std::optional<int64_t> weight_bdim;
  26336. if (weight) {
  26337. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  26338. }
  26339. std::optional<Tensor> bias_value;
  26340. std::optional<int64_t> bias_bdim;
  26341. if (bias) {
  26342. std::tie(bias_value, bias_bdim) = unwrapTensorAtLevel(bias.value(), cur_level);
  26343. }
  26344. auto results = batch_rule(input_value, input_bdim, weight_value, weight_bdim, bias_value, bias_bdim, running_mean_value, running_mean_bdim, running_var_value, running_var_bdim, momentum, eps);
  26345. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level), makeBatched(std::get<6>(results), std::get<7>(results), cur_level), makeBatched(std::get<8>(results), std::get<9>(results), cur_level), makeBatched(std::get<10>(results), std::get<11>(results), cur_level));
  26346. }
  26347. template <typename batch_rule_t, batch_rule_t batch_rule>
  26348. at::Tensor resize_as_generated_plumbing(const at::Tensor & self, const at::Tensor & the_template, ::std::optional<at::MemoryFormat> memory_format) {
  26349. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26350. auto maybe_layer = maybeCurrentDynamicLayer();
  26351. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  26352. int64_t cur_level = maybe_layer->layerId();
  26353. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(the_template, cur_level)) {
  26354. return at::_ops::resize_as::call(self, the_template, memory_format);
  26355. }
  26356. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  26357. auto [the_template_value, the_template_bdim] = unwrapTensorAtLevel(the_template, cur_level);
  26358. auto results = batch_rule(self_value, self_bdim, the_template_value, the_template_bdim, memory_format);
  26359. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  26360. }
  26361. template <typename batch_rule_t, batch_rule_t batch_rule>
  26362. at::Tensor resize_as_sparse_generated_plumbing(const at::Tensor & self, const at::Tensor & the_template) {
  26363. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26364. auto maybe_layer = maybeCurrentDynamicLayer();
  26365. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  26366. int64_t cur_level = maybe_layer->layerId();
  26367. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(the_template, cur_level)) {
  26368. return at::_ops::resize_as_sparse::call(self, the_template);
  26369. }
  26370. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  26371. auto [the_template_value, the_template_bdim] = unwrapTensorAtLevel(the_template, cur_level);
  26372. auto results = batch_rule(self_value, self_bdim, the_template_value, the_template_bdim);
  26373. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  26374. }
  26375. template <typename batch_rule_t, batch_rule_t batch_rule>
  26376. at::Tensor zero_generated_plumbing(const at::Tensor & self) {
  26377. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26378. auto maybe_layer = maybeCurrentDynamicLayer();
  26379. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  26380. int64_t cur_level = maybe_layer->layerId();
  26381. if (!isBatchedAtLevel(self, cur_level)) {
  26382. return at::_ops::zero::call(self);
  26383. }
  26384. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  26385. auto results = batch_rule(self_value, self_bdim);
  26386. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  26387. }
  26388. template <typename batch_rule_t, batch_rule_t batch_rule>
  26389. at::Tensor sparse_resize_generated_plumbing(const at::Tensor & self, at::IntArrayRef size, int64_t sparse_dim, int64_t dense_dim) {
  26390. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26391. auto maybe_layer = maybeCurrentDynamicLayer();
  26392. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  26393. int64_t cur_level = maybe_layer->layerId();
  26394. if (!isBatchedAtLevel(self, cur_level)) {
  26395. return at::_ops::sparse_resize::call(self, size, sparse_dim, dense_dim);
  26396. }
  26397. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  26398. auto results = batch_rule(self_value, self_bdim, size, sparse_dim, dense_dim);
  26399. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  26400. }
  26401. template <typename batch_rule_t, batch_rule_t batch_rule>
  26402. at::Tensor sparse_resize_and_clear_generated_plumbing(const at::Tensor & self, at::IntArrayRef size, int64_t sparse_dim, int64_t dense_dim) {
  26403. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26404. auto maybe_layer = maybeCurrentDynamicLayer();
  26405. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  26406. int64_t cur_level = maybe_layer->layerId();
  26407. if (!isBatchedAtLevel(self, cur_level)) {
  26408. return at::_ops::sparse_resize_and_clear::call(self, size, sparse_dim, dense_dim);
  26409. }
  26410. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  26411. auto results = batch_rule(self_value, self_bdim, size, sparse_dim, dense_dim);
  26412. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  26413. }
  26414. template <typename batch_rule_t, batch_rule_t batch_rule>
  26415. at::Tensor _coalesced_generated_plumbing(const at::Tensor & self, bool coalesced) {
  26416. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26417. auto maybe_layer = maybeCurrentDynamicLayer();
  26418. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  26419. int64_t cur_level = maybe_layer->layerId();
  26420. if (!isBatchedAtLevel(self, cur_level)) {
  26421. return at::_ops::_coalesced::call(self, coalesced);
  26422. }
  26423. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  26424. auto results = batch_rule(self_value, self_bdim, coalesced);
  26425. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  26426. }
  26427. template <typename batch_rule_t, batch_rule_t batch_rule>
  26428. at::Tensor copy_sparse_to_sparse_generated_plumbing(const at::Tensor & self, const at::Tensor & src, bool non_blocking) {
  26429. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26430. auto maybe_layer = maybeCurrentDynamicLayer();
  26431. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  26432. int64_t cur_level = maybe_layer->layerId();
  26433. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(src, cur_level)) {
  26434. return at::_ops::copy_sparse_to_sparse::call(self, src, non_blocking);
  26435. }
  26436. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  26437. auto [src_value, src_bdim] = unwrapTensorAtLevel(src, cur_level);
  26438. auto results = batch_rule(self_value, self_bdim, src_value, src_bdim, non_blocking);
  26439. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  26440. }
  26441. template <typename batch_rule_t, batch_rule_t batch_rule>
  26442. void quantize_per_tensor_tensors_out_generated_plumbing(at::TensorList tensors, const at::Tensor & scales, const at::Tensor & zero_points, at::ScalarType dtype, at::TensorList out) {
  26443. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26444. auto maybe_layer = maybeCurrentDynamicLayer();
  26445. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  26446. int64_t cur_level = maybe_layer->layerId();
  26447. if (!isBatchedAtLevel(tensors, cur_level) && !isBatchedAtLevel(scales, cur_level) && !isBatchedAtLevel(zero_points, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  26448. return at::_ops::quantize_per_tensor_tensors_out::call(tensors, scales, zero_points, dtype, out);
  26449. }
  26450. auto [scales_value, scales_bdim] = unwrapTensorAtLevel(scales, cur_level);
  26451. auto [zero_points_value, zero_points_bdim] = unwrapTensorAtLevel(zero_points, cur_level);
  26452. batch_rule(tensors, scales_value, scales_bdim, zero_points_value, zero_points_bdim, dtype, out);
  26453. }
  26454. template <typename batch_rule_t, batch_rule_t batch_rule>
  26455. void dequantize_tensors_out_generated_plumbing(at::TensorList tensors, at::TensorList out) {
  26456. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26457. auto maybe_layer = maybeCurrentDynamicLayer();
  26458. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  26459. int64_t cur_level = maybe_layer->layerId();
  26460. if (!isBatchedAtLevel(tensors, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  26461. return at::_ops::dequantize_tensors_out::call(tensors, out);
  26462. }
  26463. batch_rule(tensors, out);
  26464. }
  26465. template <typename batch_rule_t, batch_rule_t batch_rule>
  26466. ::std::tuple<at::Tensor,at::Tensor,at::Tensor,at::Tensor,at::Tensor,at::Tensor> _fused_moving_avg_obs_fq_helper_functional_generated_plumbing(const at::Tensor & self, const at::Tensor & observer_on, const at::Tensor & fake_quant_on, const at::Tensor & running_min, const at::Tensor & running_max, const at::Tensor & scale, const at::Tensor & zero_point, double averaging_const, int64_t quant_min, int64_t quant_max, int64_t ch_axis, bool per_row_fake_quant, bool symmetric_quant) {
  26467. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26468. auto maybe_layer = maybeCurrentDynamicLayer();
  26469. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  26470. int64_t cur_level = maybe_layer->layerId();
  26471. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(observer_on, cur_level) && !isBatchedAtLevel(fake_quant_on, cur_level) && !isBatchedAtLevel(running_min, cur_level) && !isBatchedAtLevel(running_max, cur_level) && !isBatchedAtLevel(scale, cur_level) && !isBatchedAtLevel(zero_point, cur_level)) {
  26472. return at::_ops::_fused_moving_avg_obs_fq_helper_functional::call(self, observer_on, fake_quant_on, running_min, running_max, scale, zero_point, averaging_const, quant_min, quant_max, ch_axis, per_row_fake_quant, symmetric_quant);
  26473. }
  26474. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  26475. auto [observer_on_value, observer_on_bdim] = unwrapTensorAtLevel(observer_on, cur_level);
  26476. auto [fake_quant_on_value, fake_quant_on_bdim] = unwrapTensorAtLevel(fake_quant_on, cur_level);
  26477. auto [running_min_value, running_min_bdim] = unwrapTensorAtLevel(running_min, cur_level);
  26478. auto [running_max_value, running_max_bdim] = unwrapTensorAtLevel(running_max, cur_level);
  26479. auto [scale_value, scale_bdim] = unwrapTensorAtLevel(scale, cur_level);
  26480. auto [zero_point_value, zero_point_bdim] = unwrapTensorAtLevel(zero_point, cur_level);
  26481. auto results = batch_rule(self_value, self_bdim, observer_on_value, observer_on_bdim, fake_quant_on_value, fake_quant_on_bdim, running_min_value, running_min_bdim, running_max_value, running_max_bdim, scale_value, scale_bdim, zero_point_value, zero_point_bdim, averaging_const, quant_min, quant_max, ch_axis, per_row_fake_quant, symmetric_quant);
  26482. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level), makeBatched(std::get<4>(results), std::get<5>(results), cur_level), makeBatched(std::get<6>(results), std::get<7>(results), cur_level), makeBatched(std::get<8>(results), std::get<9>(results), cur_level), makeBatched(std::get<10>(results), std::get<11>(results), cur_level));
  26483. }
  26484. template <typename batch_rule_t, batch_rule_t batch_rule>
  26485. void lstm_mps_backward_out_generated_plumbing(const ::std::optional<at::Tensor> & grad_y, const ::std::optional<at::Tensor> & grad_hy, const ::std::optional<at::Tensor> & grad_cy, const at::Tensor & z_state, const at::Tensor & cell_state_fwd, const at::Tensor & input, const at::Tensor & layersOutputs, at::TensorList hx, at::TensorList params, bool has_biases, int64_t num_layers, double dropout, bool train, bool bidirectional, bool batch_first, at::Tensor & out0, at::TensorList out1, at::TensorList out2) {
  26486. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26487. auto maybe_layer = maybeCurrentDynamicLayer();
  26488. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  26489. int64_t cur_level = maybe_layer->layerId();
  26490. if (!isBatchedAtLevel(grad_y, cur_level) && !isBatchedAtLevel(grad_hy, cur_level) && !isBatchedAtLevel(grad_cy, cur_level) && !isBatchedAtLevel(z_state, cur_level) && !isBatchedAtLevel(cell_state_fwd, cur_level) && !isBatchedAtLevel(input, cur_level) && !isBatchedAtLevel(layersOutputs, cur_level) && !isBatchedAtLevel(hx, cur_level) && !isBatchedAtLevel(params, cur_level) && !isBatchedAtLevel(out0, cur_level) && !isBatchedAtLevel(out1, cur_level) && !isBatchedAtLevel(out2, cur_level)) {
  26491. return at::_ops::lstm_mps_backward_out::call(grad_y, grad_hy, grad_cy, z_state, cell_state_fwd, input, layersOutputs, hx, params, has_biases, num_layers, dropout, train, bidirectional, batch_first, out0, out1, out2);
  26492. }
  26493. auto [z_state_value, z_state_bdim] = unwrapTensorAtLevel(z_state, cur_level);
  26494. auto [cell_state_fwd_value, cell_state_fwd_bdim] = unwrapTensorAtLevel(cell_state_fwd, cur_level);
  26495. auto [input_value, input_bdim] = unwrapTensorAtLevel(input, cur_level);
  26496. auto [layersOutputs_value, layersOutputs_bdim] = unwrapTensorAtLevel(layersOutputs, cur_level);
  26497. auto [out0_value, out0_bdim] = unwrapTensorAtLevel(out0, cur_level);
  26498. std::optional<Tensor> grad_y_value;
  26499. std::optional<int64_t> grad_y_bdim;
  26500. if (grad_y) {
  26501. std::tie(grad_y_value, grad_y_bdim) = unwrapTensorAtLevel(grad_y.value(), cur_level);
  26502. }
  26503. std::optional<Tensor> grad_hy_value;
  26504. std::optional<int64_t> grad_hy_bdim;
  26505. if (grad_hy) {
  26506. std::tie(grad_hy_value, grad_hy_bdim) = unwrapTensorAtLevel(grad_hy.value(), cur_level);
  26507. }
  26508. std::optional<Tensor> grad_cy_value;
  26509. std::optional<int64_t> grad_cy_bdim;
  26510. if (grad_cy) {
  26511. std::tie(grad_cy_value, grad_cy_bdim) = unwrapTensorAtLevel(grad_cy.value(), cur_level);
  26512. }
  26513. batch_rule(grad_y_value, grad_y_bdim, grad_hy_value, grad_hy_bdim, grad_cy_value, grad_cy_bdim, z_state_value, z_state_bdim, cell_state_fwd_value, cell_state_fwd_bdim, input_value, input_bdim, layersOutputs_value, layersOutputs_bdim, hx, params, has_biases, num_layers, dropout, train, bidirectional, batch_first, out0_value, out0_bdim, out1, out2);
  26514. }
  26515. template <typename batch_rule_t, batch_rule_t batch_rule>
  26516. at::Tensor set_source_Storage_generated_plumbing(const at::Tensor & self, at::Storage source) {
  26517. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26518. auto maybe_layer = maybeCurrentDynamicLayer();
  26519. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  26520. int64_t cur_level = maybe_layer->layerId();
  26521. if (!isBatchedAtLevel(self, cur_level)) {
  26522. return at::_ops::set_source_Storage::call(self, source);
  26523. }
  26524. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  26525. auto results = batch_rule(self_value, self_bdim, source);
  26526. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  26527. }
  26528. template <typename batch_rule_t, batch_rule_t batch_rule>
  26529. at::Tensor set_source_Storage_storage_offset_generated_plumbing(const at::Tensor & self, at::Storage source, c10::SymInt storage_offset, c10::SymIntArrayRef size, c10::SymIntArrayRef stride) {
  26530. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26531. auto maybe_layer = maybeCurrentDynamicLayer();
  26532. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  26533. int64_t cur_level = maybe_layer->layerId();
  26534. if (!isBatchedAtLevel(self, cur_level)) {
  26535. return at::_ops::set_source_Storage_storage_offset::call(self, source, storage_offset, size, stride);
  26536. }
  26537. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  26538. auto results = batch_rule(self_value, self_bdim, source, storage_offset, size, stride);
  26539. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  26540. }
  26541. template <typename batch_rule_t, batch_rule_t batch_rule>
  26542. at::Tensor set_source_Tensor_generated_plumbing(const at::Tensor & self, const at::Tensor & source) {
  26543. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26544. auto maybe_layer = maybeCurrentDynamicLayer();
  26545. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  26546. int64_t cur_level = maybe_layer->layerId();
  26547. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(source, cur_level)) {
  26548. return at::_ops::set_source_Tensor::call(self, source);
  26549. }
  26550. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  26551. auto [source_value, source_bdim] = unwrapTensorAtLevel(source, cur_level);
  26552. auto results = batch_rule(self_value, self_bdim, source_value, source_bdim);
  26553. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  26554. }
  26555. template <typename batch_rule_t, batch_rule_t batch_rule>
  26556. at::Tensor set_generated_plumbing(const at::Tensor & self) {
  26557. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26558. auto maybe_layer = maybeCurrentDynamicLayer();
  26559. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  26560. int64_t cur_level = maybe_layer->layerId();
  26561. if (!isBatchedAtLevel(self, cur_level)) {
  26562. return at::_ops::set::call(self);
  26563. }
  26564. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  26565. auto results = batch_rule(self_value, self_bdim);
  26566. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  26567. }
  26568. template <typename batch_rule_t, batch_rule_t batch_rule>
  26569. at::Tensor random_from_generated_plumbing(const at::Tensor & self, int64_t from, ::std::optional<int64_t> to, ::std::optional<at::Generator> generator) {
  26570. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26571. auto maybe_layer = maybeCurrentDynamicLayer();
  26572. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  26573. int64_t cur_level = maybe_layer->layerId();
  26574. if (!isBatchedAtLevel(self, cur_level)) {
  26575. return at::_ops::random_from::call(self, from, to, generator);
  26576. }
  26577. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  26578. auto results = batch_rule(self_value, self_bdim, from, to, generator);
  26579. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  26580. }
  26581. template <typename batch_rule_t, batch_rule_t batch_rule>
  26582. at::Tensor random_to_generated_plumbing(const at::Tensor & self, int64_t to, ::std::optional<at::Generator> generator) {
  26583. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26584. auto maybe_layer = maybeCurrentDynamicLayer();
  26585. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  26586. int64_t cur_level = maybe_layer->layerId();
  26587. if (!isBatchedAtLevel(self, cur_level)) {
  26588. return at::_ops::random_to::call(self, to, generator);
  26589. }
  26590. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  26591. auto results = batch_rule(self_value, self_bdim, to, generator);
  26592. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  26593. }
  26594. template <typename batch_rule_t, batch_rule_t batch_rule>
  26595. at::Tensor random_generated_plumbing(const at::Tensor & self, ::std::optional<at::Generator> generator) {
  26596. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26597. auto maybe_layer = maybeCurrentDynamicLayer();
  26598. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  26599. int64_t cur_level = maybe_layer->layerId();
  26600. if (!isBatchedAtLevel(self, cur_level)) {
  26601. return at::_ops::random::call(self, generator);
  26602. }
  26603. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  26604. auto results = batch_rule(self_value, self_bdim, generator);
  26605. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  26606. }
  26607. template <typename batch_rule_t, batch_rule_t batch_rule>
  26608. at::Tensor uniform_generated_plumbing(const at::Tensor & self, double from, double to, ::std::optional<at::Generator> generator) {
  26609. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26610. auto maybe_layer = maybeCurrentDynamicLayer();
  26611. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  26612. int64_t cur_level = maybe_layer->layerId();
  26613. if (!isBatchedAtLevel(self, cur_level)) {
  26614. return at::_ops::uniform::call(self, from, to, generator);
  26615. }
  26616. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  26617. auto results = batch_rule(self_value, self_bdim, from, to, generator);
  26618. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  26619. }
  26620. template <typename batch_rule_t, batch_rule_t batch_rule>
  26621. at::Tensor cauchy_generated_plumbing(const at::Tensor & self, double median, double sigma, ::std::optional<at::Generator> generator) {
  26622. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26623. auto maybe_layer = maybeCurrentDynamicLayer();
  26624. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  26625. int64_t cur_level = maybe_layer->layerId();
  26626. if (!isBatchedAtLevel(self, cur_level)) {
  26627. return at::_ops::cauchy::call(self, median, sigma, generator);
  26628. }
  26629. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  26630. auto results = batch_rule(self_value, self_bdim, median, sigma, generator);
  26631. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  26632. }
  26633. template <typename batch_rule_t, batch_rule_t batch_rule>
  26634. at::Tensor log_normal_generated_plumbing(const at::Tensor & self, double mean, double std, ::std::optional<at::Generator> generator) {
  26635. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26636. auto maybe_layer = maybeCurrentDynamicLayer();
  26637. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  26638. int64_t cur_level = maybe_layer->layerId();
  26639. if (!isBatchedAtLevel(self, cur_level)) {
  26640. return at::_ops::log_normal::call(self, mean, std, generator);
  26641. }
  26642. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  26643. auto results = batch_rule(self_value, self_bdim, mean, std, generator);
  26644. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  26645. }
  26646. template <typename batch_rule_t, batch_rule_t batch_rule>
  26647. at::Tensor exponential_generated_plumbing(const at::Tensor & self, double lambd, ::std::optional<at::Generator> generator) {
  26648. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26649. auto maybe_layer = maybeCurrentDynamicLayer();
  26650. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  26651. int64_t cur_level = maybe_layer->layerId();
  26652. if (!isBatchedAtLevel(self, cur_level)) {
  26653. return at::_ops::exponential::call(self, lambd, generator);
  26654. }
  26655. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  26656. auto results = batch_rule(self_value, self_bdim, lambd, generator);
  26657. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  26658. }
  26659. template <typename batch_rule_t, batch_rule_t batch_rule>
  26660. at::Tensor geometric_generated_plumbing(const at::Tensor & self, double p, ::std::optional<at::Generator> generator) {
  26661. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26662. auto maybe_layer = maybeCurrentDynamicLayer();
  26663. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  26664. int64_t cur_level = maybe_layer->layerId();
  26665. if (!isBatchedAtLevel(self, cur_level)) {
  26666. return at::_ops::geometric::call(self, p, generator);
  26667. }
  26668. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  26669. auto results = batch_rule(self_value, self_bdim, p, generator);
  26670. return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
  26671. }
  26672. template <typename batch_rule_t, batch_rule_t batch_rule>
  26673. void _histogramdd_bin_edges_out_generated_plumbing(const at::Tensor & self, at::IntArrayRef bins, ::std::optional<at::ArrayRef<double>> range, const ::std::optional<at::Tensor> & weight, bool density, at::TensorList out) {
  26674. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26675. auto maybe_layer = maybeCurrentDynamicLayer();
  26676. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  26677. int64_t cur_level = maybe_layer->layerId();
  26678. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(weight, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  26679. return at::_ops::_histogramdd_bin_edges_out::call(self, bins, range, weight, density, out);
  26680. }
  26681. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  26682. std::optional<Tensor> weight_value;
  26683. std::optional<int64_t> weight_bdim;
  26684. if (weight) {
  26685. std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(weight.value(), cur_level);
  26686. }
  26687. batch_rule(self_value, self_bdim, bins, range, weight_value, weight_bdim, density, out);
  26688. }
  26689. template <typename batch_rule_t, batch_rule_t batch_rule>
  26690. void _amp_foreach_non_finite_check_and_unscale_out_generated_plumbing(at::TensorList self, at::Tensor & found_inf, const at::Tensor & inv_scale, at::TensorList out) {
  26691. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26692. auto maybe_layer = maybeCurrentDynamicLayer();
  26693. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  26694. int64_t cur_level = maybe_layer->layerId();
  26695. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(found_inf, cur_level) && !isBatchedAtLevel(inv_scale, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  26696. return at::_ops::_amp_foreach_non_finite_check_and_unscale_out::call(self, found_inf, inv_scale, out);
  26697. }
  26698. auto [found_inf_value, found_inf_bdim] = unwrapTensorAtLevel(found_inf, cur_level);
  26699. auto [inv_scale_value, inv_scale_bdim] = unwrapTensorAtLevel(inv_scale, cur_level);
  26700. batch_rule(self, found_inf_value, found_inf_bdim, inv_scale_value, inv_scale_bdim, out);
  26701. }
  26702. template <typename batch_rule_t, batch_rule_t batch_rule>
  26703. ::std::tuple<::std::vector<at::Tensor>,at::Tensor> _amp_foreach_non_finite_check_and_unscale_generated_plumbing(at::TensorList self, const at::Tensor & found_inf, const at::Tensor & inv_scale) {
  26704. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26705. auto maybe_layer = maybeCurrentDynamicLayer();
  26706. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  26707. int64_t cur_level = maybe_layer->layerId();
  26708. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(found_inf, cur_level) && !isBatchedAtLevel(inv_scale, cur_level)) {
  26709. return at::_ops::_amp_foreach_non_finite_check_and_unscale::call(self, found_inf, inv_scale);
  26710. }
  26711. auto [found_inf_value, found_inf_bdim] = unwrapTensorAtLevel(found_inf, cur_level);
  26712. auto [inv_scale_value, inv_scale_bdim] = unwrapTensorAtLevel(inv_scale, cur_level);
  26713. auto results = batch_rule(self, found_inf_value, found_inf_bdim, inv_scale_value, inv_scale_bdim);
  26714. return std::make_tuple(makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  26715. }
  26716. template <typename batch_rule_t, batch_rule_t batch_rule>
  26717. ::std::tuple<at::Tensor,at::Tensor> _amp_update_scale_generated_plumbing(const at::Tensor & self, const at::Tensor & growth_tracker, const at::Tensor & found_inf, double scale_growth_factor, double scale_backoff_factor, int64_t growth_interval) {
  26718. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26719. auto maybe_layer = maybeCurrentDynamicLayer();
  26720. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  26721. int64_t cur_level = maybe_layer->layerId();
  26722. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(growth_tracker, cur_level) && !isBatchedAtLevel(found_inf, cur_level)) {
  26723. return at::_ops::_amp_update_scale::call(self, growth_tracker, found_inf, scale_growth_factor, scale_backoff_factor, growth_interval);
  26724. }
  26725. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  26726. auto [growth_tracker_value, growth_tracker_bdim] = unwrapTensorAtLevel(growth_tracker, cur_level);
  26727. auto [found_inf_value, found_inf_bdim] = unwrapTensorAtLevel(found_inf, cur_level);
  26728. auto results = batch_rule(self_value, self_bdim, growth_tracker_value, growth_tracker_bdim, found_inf_value, found_inf_bdim, scale_growth_factor, scale_backoff_factor, growth_interval);
  26729. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  26730. }
  26731. template <typename batch_rule_t, batch_rule_t batch_rule>
  26732. void _foreach_add_Scalar_out_generated_plumbing(at::TensorList self, const at::Scalar & scalar, at::TensorList out) {
  26733. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26734. auto maybe_layer = maybeCurrentDynamicLayer();
  26735. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  26736. int64_t cur_level = maybe_layer->layerId();
  26737. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  26738. return at::_ops::_foreach_add_Scalar_out::call(self, scalar, out);
  26739. }
  26740. batch_rule(self, scalar, out);
  26741. }
  26742. template <typename batch_rule_t, batch_rule_t batch_rule>
  26743. void _foreach_add_List_out_generated_plumbing(at::TensorList self, at::TensorList other, const at::Scalar & alpha, at::TensorList out) {
  26744. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26745. auto maybe_layer = maybeCurrentDynamicLayer();
  26746. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  26747. int64_t cur_level = maybe_layer->layerId();
  26748. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  26749. return at::_ops::_foreach_add_List_out::call(self, other, alpha, out);
  26750. }
  26751. batch_rule(self, other, alpha, out);
  26752. }
  26753. template <typename batch_rule_t, batch_rule_t batch_rule>
  26754. void _foreach_add_ScalarList_out_generated_plumbing(at::TensorList self, at::ArrayRef<at::Scalar> scalars, at::TensorList out) {
  26755. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26756. auto maybe_layer = maybeCurrentDynamicLayer();
  26757. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  26758. int64_t cur_level = maybe_layer->layerId();
  26759. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  26760. return at::_ops::_foreach_add_ScalarList_out::call(self, scalars, out);
  26761. }
  26762. batch_rule(self, scalars, out);
  26763. }
  26764. template <typename batch_rule_t, batch_rule_t batch_rule>
  26765. void _foreach_add_Tensor_out_generated_plumbing(at::TensorList self, const at::Tensor & other, const at::Scalar & alpha, at::TensorList out) {
  26766. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26767. auto maybe_layer = maybeCurrentDynamicLayer();
  26768. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  26769. int64_t cur_level = maybe_layer->layerId();
  26770. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  26771. return at::_ops::_foreach_add_Tensor_out::call(self, other, alpha, out);
  26772. }
  26773. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  26774. batch_rule(self, other_value, other_bdim, alpha, out);
  26775. }
  26776. template <typename batch_rule_t, batch_rule_t batch_rule>
  26777. void _foreach_sub_Scalar_out_generated_plumbing(at::TensorList self, const at::Scalar & scalar, at::TensorList out) {
  26778. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26779. auto maybe_layer = maybeCurrentDynamicLayer();
  26780. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  26781. int64_t cur_level = maybe_layer->layerId();
  26782. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  26783. return at::_ops::_foreach_sub_Scalar_out::call(self, scalar, out);
  26784. }
  26785. batch_rule(self, scalar, out);
  26786. }
  26787. template <typename batch_rule_t, batch_rule_t batch_rule>
  26788. void _foreach_sub_List_out_generated_plumbing(at::TensorList self, at::TensorList other, const at::Scalar & alpha, at::TensorList out) {
  26789. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26790. auto maybe_layer = maybeCurrentDynamicLayer();
  26791. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  26792. int64_t cur_level = maybe_layer->layerId();
  26793. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  26794. return at::_ops::_foreach_sub_List_out::call(self, other, alpha, out);
  26795. }
  26796. batch_rule(self, other, alpha, out);
  26797. }
  26798. template <typename batch_rule_t, batch_rule_t batch_rule>
  26799. void _foreach_sub_ScalarList_out_generated_plumbing(at::TensorList self, at::ArrayRef<at::Scalar> scalars, at::TensorList out) {
  26800. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26801. auto maybe_layer = maybeCurrentDynamicLayer();
  26802. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  26803. int64_t cur_level = maybe_layer->layerId();
  26804. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  26805. return at::_ops::_foreach_sub_ScalarList_out::call(self, scalars, out);
  26806. }
  26807. batch_rule(self, scalars, out);
  26808. }
  26809. template <typename batch_rule_t, batch_rule_t batch_rule>
  26810. void _foreach_mul_Scalar_out_generated_plumbing(at::TensorList self, const at::Scalar & scalar, at::TensorList out) {
  26811. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26812. auto maybe_layer = maybeCurrentDynamicLayer();
  26813. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  26814. int64_t cur_level = maybe_layer->layerId();
  26815. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  26816. return at::_ops::_foreach_mul_Scalar_out::call(self, scalar, out);
  26817. }
  26818. batch_rule(self, scalar, out);
  26819. }
  26820. template <typename batch_rule_t, batch_rule_t batch_rule>
  26821. void _foreach_mul_List_out_generated_plumbing(at::TensorList self, at::TensorList other, at::TensorList out) {
  26822. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26823. auto maybe_layer = maybeCurrentDynamicLayer();
  26824. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  26825. int64_t cur_level = maybe_layer->layerId();
  26826. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  26827. return at::_ops::_foreach_mul_List_out::call(self, other, out);
  26828. }
  26829. batch_rule(self, other, out);
  26830. }
  26831. template <typename batch_rule_t, batch_rule_t batch_rule>
  26832. void _foreach_mul_ScalarList_out_generated_plumbing(at::TensorList self, at::ArrayRef<at::Scalar> scalars, at::TensorList out) {
  26833. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26834. auto maybe_layer = maybeCurrentDynamicLayer();
  26835. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  26836. int64_t cur_level = maybe_layer->layerId();
  26837. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  26838. return at::_ops::_foreach_mul_ScalarList_out::call(self, scalars, out);
  26839. }
  26840. batch_rule(self, scalars, out);
  26841. }
  26842. template <typename batch_rule_t, batch_rule_t batch_rule>
  26843. void _foreach_mul_Tensor_out_generated_plumbing(at::TensorList self, const at::Tensor & other, at::TensorList out) {
  26844. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26845. auto maybe_layer = maybeCurrentDynamicLayer();
  26846. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  26847. int64_t cur_level = maybe_layer->layerId();
  26848. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  26849. return at::_ops::_foreach_mul_Tensor_out::call(self, other, out);
  26850. }
  26851. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  26852. batch_rule(self, other_value, other_bdim, out);
  26853. }
  26854. template <typename batch_rule_t, batch_rule_t batch_rule>
  26855. void _foreach_div_Scalar_out_generated_plumbing(at::TensorList self, const at::Scalar & scalar, at::TensorList out) {
  26856. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26857. auto maybe_layer = maybeCurrentDynamicLayer();
  26858. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  26859. int64_t cur_level = maybe_layer->layerId();
  26860. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  26861. return at::_ops::_foreach_div_Scalar_out::call(self, scalar, out);
  26862. }
  26863. batch_rule(self, scalar, out);
  26864. }
  26865. template <typename batch_rule_t, batch_rule_t batch_rule>
  26866. void _foreach_div_List_out_generated_plumbing(at::TensorList self, at::TensorList other, at::TensorList out) {
  26867. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26868. auto maybe_layer = maybeCurrentDynamicLayer();
  26869. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  26870. int64_t cur_level = maybe_layer->layerId();
  26871. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  26872. return at::_ops::_foreach_div_List_out::call(self, other, out);
  26873. }
  26874. batch_rule(self, other, out);
  26875. }
  26876. template <typename batch_rule_t, batch_rule_t batch_rule>
  26877. void _foreach_div_ScalarList_out_generated_plumbing(at::TensorList self, at::ArrayRef<at::Scalar> scalars, at::TensorList out) {
  26878. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26879. auto maybe_layer = maybeCurrentDynamicLayer();
  26880. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  26881. int64_t cur_level = maybe_layer->layerId();
  26882. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  26883. return at::_ops::_foreach_div_ScalarList_out::call(self, scalars, out);
  26884. }
  26885. batch_rule(self, scalars, out);
  26886. }
  26887. template <typename batch_rule_t, batch_rule_t batch_rule>
  26888. void _foreach_div_Tensor_out_generated_plumbing(at::TensorList self, const at::Tensor & other, at::TensorList out) {
  26889. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26890. auto maybe_layer = maybeCurrentDynamicLayer();
  26891. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  26892. int64_t cur_level = maybe_layer->layerId();
  26893. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  26894. return at::_ops::_foreach_div_Tensor_out::call(self, other, out);
  26895. }
  26896. auto [other_value, other_bdim] = unwrapTensorAtLevel(other, cur_level);
  26897. batch_rule(self, other_value, other_bdim, out);
  26898. }
  26899. template <typename batch_rule_t, batch_rule_t batch_rule>
  26900. void _foreach_clamp_max_Scalar_out_generated_plumbing(at::TensorList self, const at::Scalar & scalar, at::TensorList out) {
  26901. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26902. auto maybe_layer = maybeCurrentDynamicLayer();
  26903. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  26904. int64_t cur_level = maybe_layer->layerId();
  26905. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  26906. return at::_ops::_foreach_clamp_max_Scalar_out::call(self, scalar, out);
  26907. }
  26908. batch_rule(self, scalar, out);
  26909. }
  26910. template <typename batch_rule_t, batch_rule_t batch_rule>
  26911. void _foreach_clamp_max_List_out_generated_plumbing(at::TensorList self, at::TensorList other, at::TensorList out) {
  26912. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26913. auto maybe_layer = maybeCurrentDynamicLayer();
  26914. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  26915. int64_t cur_level = maybe_layer->layerId();
  26916. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  26917. return at::_ops::_foreach_clamp_max_List_out::call(self, other, out);
  26918. }
  26919. batch_rule(self, other, out);
  26920. }
  26921. template <typename batch_rule_t, batch_rule_t batch_rule>
  26922. void _foreach_clamp_max_ScalarList_out_generated_plumbing(at::TensorList self, at::ArrayRef<at::Scalar> scalars, at::TensorList out) {
  26923. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26924. auto maybe_layer = maybeCurrentDynamicLayer();
  26925. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  26926. int64_t cur_level = maybe_layer->layerId();
  26927. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  26928. return at::_ops::_foreach_clamp_max_ScalarList_out::call(self, scalars, out);
  26929. }
  26930. batch_rule(self, scalars, out);
  26931. }
  26932. template <typename batch_rule_t, batch_rule_t batch_rule>
  26933. void _foreach_clamp_min_Scalar_out_generated_plumbing(at::TensorList self, const at::Scalar & scalar, at::TensorList out) {
  26934. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26935. auto maybe_layer = maybeCurrentDynamicLayer();
  26936. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  26937. int64_t cur_level = maybe_layer->layerId();
  26938. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  26939. return at::_ops::_foreach_clamp_min_Scalar_out::call(self, scalar, out);
  26940. }
  26941. batch_rule(self, scalar, out);
  26942. }
  26943. template <typename batch_rule_t, batch_rule_t batch_rule>
  26944. void _foreach_clamp_min_List_out_generated_plumbing(at::TensorList self, at::TensorList other, at::TensorList out) {
  26945. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26946. auto maybe_layer = maybeCurrentDynamicLayer();
  26947. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  26948. int64_t cur_level = maybe_layer->layerId();
  26949. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  26950. return at::_ops::_foreach_clamp_min_List_out::call(self, other, out);
  26951. }
  26952. batch_rule(self, other, out);
  26953. }
  26954. template <typename batch_rule_t, batch_rule_t batch_rule>
  26955. void _foreach_clamp_min_ScalarList_out_generated_plumbing(at::TensorList self, at::ArrayRef<at::Scalar> scalars, at::TensorList out) {
  26956. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26957. auto maybe_layer = maybeCurrentDynamicLayer();
  26958. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  26959. int64_t cur_level = maybe_layer->layerId();
  26960. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  26961. return at::_ops::_foreach_clamp_min_ScalarList_out::call(self, scalars, out);
  26962. }
  26963. batch_rule(self, scalars, out);
  26964. }
  26965. template <typename batch_rule_t, batch_rule_t batch_rule>
  26966. void _foreach_maximum_Scalar_out_generated_plumbing(at::TensorList self, const at::Scalar & scalar, at::TensorList out) {
  26967. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26968. auto maybe_layer = maybeCurrentDynamicLayer();
  26969. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  26970. int64_t cur_level = maybe_layer->layerId();
  26971. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  26972. return at::_ops::_foreach_maximum_Scalar_out::call(self, scalar, out);
  26973. }
  26974. batch_rule(self, scalar, out);
  26975. }
  26976. template <typename batch_rule_t, batch_rule_t batch_rule>
  26977. void _foreach_maximum_List_out_generated_plumbing(at::TensorList self, at::TensorList other, at::TensorList out) {
  26978. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26979. auto maybe_layer = maybeCurrentDynamicLayer();
  26980. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  26981. int64_t cur_level = maybe_layer->layerId();
  26982. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  26983. return at::_ops::_foreach_maximum_List_out::call(self, other, out);
  26984. }
  26985. batch_rule(self, other, out);
  26986. }
  26987. template <typename batch_rule_t, batch_rule_t batch_rule>
  26988. void _foreach_maximum_ScalarList_out_generated_plumbing(at::TensorList self, at::ArrayRef<at::Scalar> scalars, at::TensorList out) {
  26989. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  26990. auto maybe_layer = maybeCurrentDynamicLayer();
  26991. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  26992. int64_t cur_level = maybe_layer->layerId();
  26993. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  26994. return at::_ops::_foreach_maximum_ScalarList_out::call(self, scalars, out);
  26995. }
  26996. batch_rule(self, scalars, out);
  26997. }
  26998. template <typename batch_rule_t, batch_rule_t batch_rule>
  26999. void _foreach_minimum_Scalar_out_generated_plumbing(at::TensorList self, const at::Scalar & scalar, at::TensorList out) {
  27000. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27001. auto maybe_layer = maybeCurrentDynamicLayer();
  27002. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27003. int64_t cur_level = maybe_layer->layerId();
  27004. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27005. return at::_ops::_foreach_minimum_Scalar_out::call(self, scalar, out);
  27006. }
  27007. batch_rule(self, scalar, out);
  27008. }
  27009. template <typename batch_rule_t, batch_rule_t batch_rule>
  27010. void _foreach_minimum_List_out_generated_plumbing(at::TensorList self, at::TensorList other, at::TensorList out) {
  27011. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27012. auto maybe_layer = maybeCurrentDynamicLayer();
  27013. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27014. int64_t cur_level = maybe_layer->layerId();
  27015. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(other, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27016. return at::_ops::_foreach_minimum_List_out::call(self, other, out);
  27017. }
  27018. batch_rule(self, other, out);
  27019. }
  27020. template <typename batch_rule_t, batch_rule_t batch_rule>
  27021. void _foreach_minimum_ScalarList_out_generated_plumbing(at::TensorList self, at::ArrayRef<at::Scalar> scalars, at::TensorList out) {
  27022. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27023. auto maybe_layer = maybeCurrentDynamicLayer();
  27024. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27025. int64_t cur_level = maybe_layer->layerId();
  27026. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27027. return at::_ops::_foreach_minimum_ScalarList_out::call(self, scalars, out);
  27028. }
  27029. batch_rule(self, scalars, out);
  27030. }
  27031. template <typename batch_rule_t, batch_rule_t batch_rule>
  27032. void _foreach_addcdiv_Scalar_out_generated_plumbing(at::TensorList self, at::TensorList tensor1, at::TensorList tensor2, const at::Scalar & value, at::TensorList out) {
  27033. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27034. auto maybe_layer = maybeCurrentDynamicLayer();
  27035. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27036. int64_t cur_level = maybe_layer->layerId();
  27037. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(tensor1, cur_level) && !isBatchedAtLevel(tensor2, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27038. return at::_ops::_foreach_addcdiv_Scalar_out::call(self, tensor1, tensor2, value, out);
  27039. }
  27040. batch_rule(self, tensor1, tensor2, value, out);
  27041. }
  27042. template <typename batch_rule_t, batch_rule_t batch_rule>
  27043. void _foreach_addcdiv_ScalarList_out_generated_plumbing(at::TensorList self, at::TensorList tensor1, at::TensorList tensor2, at::ArrayRef<at::Scalar> scalars, at::TensorList out) {
  27044. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27045. auto maybe_layer = maybeCurrentDynamicLayer();
  27046. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27047. int64_t cur_level = maybe_layer->layerId();
  27048. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(tensor1, cur_level) && !isBatchedAtLevel(tensor2, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27049. return at::_ops::_foreach_addcdiv_ScalarList_out::call(self, tensor1, tensor2, scalars, out);
  27050. }
  27051. batch_rule(self, tensor1, tensor2, scalars, out);
  27052. }
  27053. template <typename batch_rule_t, batch_rule_t batch_rule>
  27054. void _foreach_addcdiv_Tensor_out_generated_plumbing(at::TensorList self, at::TensorList tensor1, at::TensorList tensor2, const at::Tensor & scalars, at::TensorList out) {
  27055. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27056. auto maybe_layer = maybeCurrentDynamicLayer();
  27057. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27058. int64_t cur_level = maybe_layer->layerId();
  27059. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(tensor1, cur_level) && !isBatchedAtLevel(tensor2, cur_level) && !isBatchedAtLevel(scalars, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27060. return at::_ops::_foreach_addcdiv_Tensor_out::call(self, tensor1, tensor2, scalars, out);
  27061. }
  27062. auto [scalars_value, scalars_bdim] = unwrapTensorAtLevel(scalars, cur_level);
  27063. batch_rule(self, tensor1, tensor2, scalars_value, scalars_bdim, out);
  27064. }
  27065. template <typename batch_rule_t, batch_rule_t batch_rule>
  27066. void _foreach_addcmul_Scalar_out_generated_plumbing(at::TensorList self, at::TensorList tensor1, at::TensorList tensor2, const at::Scalar & value, at::TensorList out) {
  27067. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27068. auto maybe_layer = maybeCurrentDynamicLayer();
  27069. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27070. int64_t cur_level = maybe_layer->layerId();
  27071. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(tensor1, cur_level) && !isBatchedAtLevel(tensor2, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27072. return at::_ops::_foreach_addcmul_Scalar_out::call(self, tensor1, tensor2, value, out);
  27073. }
  27074. batch_rule(self, tensor1, tensor2, value, out);
  27075. }
  27076. template <typename batch_rule_t, batch_rule_t batch_rule>
  27077. void _foreach_addcmul_ScalarList_out_generated_plumbing(at::TensorList self, at::TensorList tensor1, at::TensorList tensor2, at::ArrayRef<at::Scalar> scalars, at::TensorList out) {
  27078. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27079. auto maybe_layer = maybeCurrentDynamicLayer();
  27080. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27081. int64_t cur_level = maybe_layer->layerId();
  27082. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(tensor1, cur_level) && !isBatchedAtLevel(tensor2, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27083. return at::_ops::_foreach_addcmul_ScalarList_out::call(self, tensor1, tensor2, scalars, out);
  27084. }
  27085. batch_rule(self, tensor1, tensor2, scalars, out);
  27086. }
  27087. template <typename batch_rule_t, batch_rule_t batch_rule>
  27088. void _foreach_addcmul_Tensor_out_generated_plumbing(at::TensorList self, at::TensorList tensor1, at::TensorList tensor2, const at::Tensor & scalars, at::TensorList out) {
  27089. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27090. auto maybe_layer = maybeCurrentDynamicLayer();
  27091. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27092. int64_t cur_level = maybe_layer->layerId();
  27093. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(tensor1, cur_level) && !isBatchedAtLevel(tensor2, cur_level) && !isBatchedAtLevel(scalars, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27094. return at::_ops::_foreach_addcmul_Tensor_out::call(self, tensor1, tensor2, scalars, out);
  27095. }
  27096. auto [scalars_value, scalars_bdim] = unwrapTensorAtLevel(scalars, cur_level);
  27097. batch_rule(self, tensor1, tensor2, scalars_value, scalars_bdim, out);
  27098. }
  27099. template <typename batch_rule_t, batch_rule_t batch_rule>
  27100. void _foreach_abs_out_generated_plumbing(at::TensorList self, at::TensorList out) {
  27101. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27102. auto maybe_layer = maybeCurrentDynamicLayer();
  27103. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27104. int64_t cur_level = maybe_layer->layerId();
  27105. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27106. return at::_ops::_foreach_abs_out::call(self, out);
  27107. }
  27108. batch_rule(self, out);
  27109. }
  27110. template <typename batch_rule_t, batch_rule_t batch_rule>
  27111. void _foreach_acos_out_generated_plumbing(at::TensorList self, at::TensorList out) {
  27112. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27113. auto maybe_layer = maybeCurrentDynamicLayer();
  27114. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27115. int64_t cur_level = maybe_layer->layerId();
  27116. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27117. return at::_ops::_foreach_acos_out::call(self, out);
  27118. }
  27119. batch_rule(self, out);
  27120. }
  27121. template <typename batch_rule_t, batch_rule_t batch_rule>
  27122. void _foreach_asin_out_generated_plumbing(at::TensorList self, at::TensorList out) {
  27123. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27124. auto maybe_layer = maybeCurrentDynamicLayer();
  27125. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27126. int64_t cur_level = maybe_layer->layerId();
  27127. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27128. return at::_ops::_foreach_asin_out::call(self, out);
  27129. }
  27130. batch_rule(self, out);
  27131. }
  27132. template <typename batch_rule_t, batch_rule_t batch_rule>
  27133. void _foreach_atan_out_generated_plumbing(at::TensorList self, at::TensorList out) {
  27134. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27135. auto maybe_layer = maybeCurrentDynamicLayer();
  27136. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27137. int64_t cur_level = maybe_layer->layerId();
  27138. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27139. return at::_ops::_foreach_atan_out::call(self, out);
  27140. }
  27141. batch_rule(self, out);
  27142. }
  27143. template <typename batch_rule_t, batch_rule_t batch_rule>
  27144. void _foreach_ceil_out_generated_plumbing(at::TensorList self, at::TensorList out) {
  27145. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27146. auto maybe_layer = maybeCurrentDynamicLayer();
  27147. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27148. int64_t cur_level = maybe_layer->layerId();
  27149. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27150. return at::_ops::_foreach_ceil_out::call(self, out);
  27151. }
  27152. batch_rule(self, out);
  27153. }
  27154. template <typename batch_rule_t, batch_rule_t batch_rule>
  27155. void _foreach_cos_out_generated_plumbing(at::TensorList self, at::TensorList out) {
  27156. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27157. auto maybe_layer = maybeCurrentDynamicLayer();
  27158. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27159. int64_t cur_level = maybe_layer->layerId();
  27160. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27161. return at::_ops::_foreach_cos_out::call(self, out);
  27162. }
  27163. batch_rule(self, out);
  27164. }
  27165. template <typename batch_rule_t, batch_rule_t batch_rule>
  27166. void _foreach_cosh_out_generated_plumbing(at::TensorList self, at::TensorList out) {
  27167. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27168. auto maybe_layer = maybeCurrentDynamicLayer();
  27169. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27170. int64_t cur_level = maybe_layer->layerId();
  27171. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27172. return at::_ops::_foreach_cosh_out::call(self, out);
  27173. }
  27174. batch_rule(self, out);
  27175. }
  27176. template <typename batch_rule_t, batch_rule_t batch_rule>
  27177. void _foreach_erf_out_generated_plumbing(at::TensorList self, at::TensorList out) {
  27178. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27179. auto maybe_layer = maybeCurrentDynamicLayer();
  27180. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27181. int64_t cur_level = maybe_layer->layerId();
  27182. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27183. return at::_ops::_foreach_erf_out::call(self, out);
  27184. }
  27185. batch_rule(self, out);
  27186. }
  27187. template <typename batch_rule_t, batch_rule_t batch_rule>
  27188. void _foreach_erfc_out_generated_plumbing(at::TensorList self, at::TensorList out) {
  27189. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27190. auto maybe_layer = maybeCurrentDynamicLayer();
  27191. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27192. int64_t cur_level = maybe_layer->layerId();
  27193. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27194. return at::_ops::_foreach_erfc_out::call(self, out);
  27195. }
  27196. batch_rule(self, out);
  27197. }
  27198. template <typename batch_rule_t, batch_rule_t batch_rule>
  27199. void _foreach_exp_out_generated_plumbing(at::TensorList self, at::TensorList out) {
  27200. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27201. auto maybe_layer = maybeCurrentDynamicLayer();
  27202. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27203. int64_t cur_level = maybe_layer->layerId();
  27204. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27205. return at::_ops::_foreach_exp_out::call(self, out);
  27206. }
  27207. batch_rule(self, out);
  27208. }
  27209. template <typename batch_rule_t, batch_rule_t batch_rule>
  27210. void _foreach_expm1_out_generated_plumbing(at::TensorList self, at::TensorList out) {
  27211. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27212. auto maybe_layer = maybeCurrentDynamicLayer();
  27213. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27214. int64_t cur_level = maybe_layer->layerId();
  27215. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27216. return at::_ops::_foreach_expm1_out::call(self, out);
  27217. }
  27218. batch_rule(self, out);
  27219. }
  27220. template <typename batch_rule_t, batch_rule_t batch_rule>
  27221. void _foreach_floor_out_generated_plumbing(at::TensorList self, at::TensorList out) {
  27222. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27223. auto maybe_layer = maybeCurrentDynamicLayer();
  27224. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27225. int64_t cur_level = maybe_layer->layerId();
  27226. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27227. return at::_ops::_foreach_floor_out::call(self, out);
  27228. }
  27229. batch_rule(self, out);
  27230. }
  27231. template <typename batch_rule_t, batch_rule_t batch_rule>
  27232. void _foreach_frac_out_generated_plumbing(at::TensorList self, at::TensorList out) {
  27233. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27234. auto maybe_layer = maybeCurrentDynamicLayer();
  27235. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27236. int64_t cur_level = maybe_layer->layerId();
  27237. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27238. return at::_ops::_foreach_frac_out::call(self, out);
  27239. }
  27240. batch_rule(self, out);
  27241. }
  27242. template <typename batch_rule_t, batch_rule_t batch_rule>
  27243. void _foreach_lerp_List_out_generated_plumbing(at::TensorList self, at::TensorList tensors1, at::TensorList weights, at::TensorList out) {
  27244. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27245. auto maybe_layer = maybeCurrentDynamicLayer();
  27246. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27247. int64_t cur_level = maybe_layer->layerId();
  27248. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(tensors1, cur_level) && !isBatchedAtLevel(weights, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27249. return at::_ops::_foreach_lerp_List_out::call(self, tensors1, weights, out);
  27250. }
  27251. batch_rule(self, tensors1, weights, out);
  27252. }
  27253. template <typename batch_rule_t, batch_rule_t batch_rule>
  27254. void _foreach_lerp_Scalar_out_generated_plumbing(at::TensorList self, at::TensorList tensors1, const at::Scalar & weight, at::TensorList out) {
  27255. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27256. auto maybe_layer = maybeCurrentDynamicLayer();
  27257. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27258. int64_t cur_level = maybe_layer->layerId();
  27259. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(tensors1, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27260. return at::_ops::_foreach_lerp_Scalar_out::call(self, tensors1, weight, out);
  27261. }
  27262. batch_rule(self, tensors1, weight, out);
  27263. }
  27264. template <typename batch_rule_t, batch_rule_t batch_rule>
  27265. void _foreach_lerp_ScalarList_out_generated_plumbing(at::TensorList self, at::TensorList tensors1, at::ArrayRef<at::Scalar> weight, at::TensorList out) {
  27266. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27267. auto maybe_layer = maybeCurrentDynamicLayer();
  27268. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27269. int64_t cur_level = maybe_layer->layerId();
  27270. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(tensors1, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27271. return at::_ops::_foreach_lerp_ScalarList_out::call(self, tensors1, weight, out);
  27272. }
  27273. batch_rule(self, tensors1, weight, out);
  27274. }
  27275. template <typename batch_rule_t, batch_rule_t batch_rule>
  27276. void _foreach_lgamma_out_generated_plumbing(at::TensorList self, at::TensorList out) {
  27277. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27278. auto maybe_layer = maybeCurrentDynamicLayer();
  27279. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27280. int64_t cur_level = maybe_layer->layerId();
  27281. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27282. return at::_ops::_foreach_lgamma_out::call(self, out);
  27283. }
  27284. batch_rule(self, out);
  27285. }
  27286. template <typename batch_rule_t, batch_rule_t batch_rule>
  27287. void _foreach_log_out_generated_plumbing(at::TensorList self, at::TensorList out) {
  27288. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27289. auto maybe_layer = maybeCurrentDynamicLayer();
  27290. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27291. int64_t cur_level = maybe_layer->layerId();
  27292. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27293. return at::_ops::_foreach_log_out::call(self, out);
  27294. }
  27295. batch_rule(self, out);
  27296. }
  27297. template <typename batch_rule_t, batch_rule_t batch_rule>
  27298. void _foreach_log10_out_generated_plumbing(at::TensorList self, at::TensorList out) {
  27299. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27300. auto maybe_layer = maybeCurrentDynamicLayer();
  27301. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27302. int64_t cur_level = maybe_layer->layerId();
  27303. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27304. return at::_ops::_foreach_log10_out::call(self, out);
  27305. }
  27306. batch_rule(self, out);
  27307. }
  27308. template <typename batch_rule_t, batch_rule_t batch_rule>
  27309. void _foreach_log1p_out_generated_plumbing(at::TensorList self, at::TensorList out) {
  27310. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27311. auto maybe_layer = maybeCurrentDynamicLayer();
  27312. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27313. int64_t cur_level = maybe_layer->layerId();
  27314. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27315. return at::_ops::_foreach_log1p_out::call(self, out);
  27316. }
  27317. batch_rule(self, out);
  27318. }
  27319. template <typename batch_rule_t, batch_rule_t batch_rule>
  27320. void _foreach_log2_out_generated_plumbing(at::TensorList self, at::TensorList out) {
  27321. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27322. auto maybe_layer = maybeCurrentDynamicLayer();
  27323. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27324. int64_t cur_level = maybe_layer->layerId();
  27325. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27326. return at::_ops::_foreach_log2_out::call(self, out);
  27327. }
  27328. batch_rule(self, out);
  27329. }
  27330. template <typename batch_rule_t, batch_rule_t batch_rule>
  27331. void _foreach_max_out_generated_plumbing(at::TensorList self, at::TensorList out) {
  27332. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27333. auto maybe_layer = maybeCurrentDynamicLayer();
  27334. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27335. int64_t cur_level = maybe_layer->layerId();
  27336. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27337. return at::_ops::_foreach_max_out::call(self, out);
  27338. }
  27339. batch_rule(self, out);
  27340. }
  27341. template <typename batch_rule_t, batch_rule_t batch_rule>
  27342. void _foreach_neg_out_generated_plumbing(at::TensorList self, at::TensorList out) {
  27343. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27344. auto maybe_layer = maybeCurrentDynamicLayer();
  27345. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27346. int64_t cur_level = maybe_layer->layerId();
  27347. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27348. return at::_ops::_foreach_neg_out::call(self, out);
  27349. }
  27350. batch_rule(self, out);
  27351. }
  27352. template <typename batch_rule_t, batch_rule_t batch_rule>
  27353. void _foreach_norm_Scalar_out_generated_plumbing(at::TensorList self, const at::Scalar & ord, ::std::optional<at::ScalarType> dtype, at::TensorList out) {
  27354. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27355. auto maybe_layer = maybeCurrentDynamicLayer();
  27356. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27357. int64_t cur_level = maybe_layer->layerId();
  27358. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27359. return at::_ops::_foreach_norm_Scalar_out::call(self, ord, dtype, out);
  27360. }
  27361. batch_rule(self, ord, dtype, out);
  27362. }
  27363. template <typename batch_rule_t, batch_rule_t batch_rule>
  27364. void _foreach_pow_List_out_generated_plumbing(at::TensorList self, at::TensorList exponent, at::TensorList out) {
  27365. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27366. auto maybe_layer = maybeCurrentDynamicLayer();
  27367. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27368. int64_t cur_level = maybe_layer->layerId();
  27369. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(exponent, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27370. return at::_ops::_foreach_pow_List_out::call(self, exponent, out);
  27371. }
  27372. batch_rule(self, exponent, out);
  27373. }
  27374. template <typename batch_rule_t, batch_rule_t batch_rule>
  27375. void _foreach_pow_Scalar_out_generated_plumbing(at::TensorList self, const at::Scalar & exponent, at::TensorList out) {
  27376. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27377. auto maybe_layer = maybeCurrentDynamicLayer();
  27378. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27379. int64_t cur_level = maybe_layer->layerId();
  27380. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27381. return at::_ops::_foreach_pow_Scalar_out::call(self, exponent, out);
  27382. }
  27383. batch_rule(self, exponent, out);
  27384. }
  27385. template <typename batch_rule_t, batch_rule_t batch_rule>
  27386. void _foreach_pow_ScalarList_out_generated_plumbing(at::TensorList self, at::ArrayRef<at::Scalar> exponent, at::TensorList out) {
  27387. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27388. auto maybe_layer = maybeCurrentDynamicLayer();
  27389. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27390. int64_t cur_level = maybe_layer->layerId();
  27391. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27392. return at::_ops::_foreach_pow_ScalarList_out::call(self, exponent, out);
  27393. }
  27394. batch_rule(self, exponent, out);
  27395. }
  27396. template <typename batch_rule_t, batch_rule_t batch_rule>
  27397. void _foreach_reciprocal_out_generated_plumbing(at::TensorList self, at::TensorList out) {
  27398. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27399. auto maybe_layer = maybeCurrentDynamicLayer();
  27400. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27401. int64_t cur_level = maybe_layer->layerId();
  27402. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27403. return at::_ops::_foreach_reciprocal_out::call(self, out);
  27404. }
  27405. batch_rule(self, out);
  27406. }
  27407. template <typename batch_rule_t, batch_rule_t batch_rule>
  27408. void _foreach_round_out_generated_plumbing(at::TensorList self, at::TensorList out) {
  27409. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27410. auto maybe_layer = maybeCurrentDynamicLayer();
  27411. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27412. int64_t cur_level = maybe_layer->layerId();
  27413. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27414. return at::_ops::_foreach_round_out::call(self, out);
  27415. }
  27416. batch_rule(self, out);
  27417. }
  27418. template <typename batch_rule_t, batch_rule_t batch_rule>
  27419. void _foreach_rsqrt_out_generated_plumbing(at::TensorList self, at::TensorList out) {
  27420. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27421. auto maybe_layer = maybeCurrentDynamicLayer();
  27422. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27423. int64_t cur_level = maybe_layer->layerId();
  27424. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27425. return at::_ops::_foreach_rsqrt_out::call(self, out);
  27426. }
  27427. batch_rule(self, out);
  27428. }
  27429. template <typename batch_rule_t, batch_rule_t batch_rule>
  27430. void _foreach_sigmoid_out_generated_plumbing(at::TensorList self, at::TensorList out) {
  27431. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27432. auto maybe_layer = maybeCurrentDynamicLayer();
  27433. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27434. int64_t cur_level = maybe_layer->layerId();
  27435. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27436. return at::_ops::_foreach_sigmoid_out::call(self, out);
  27437. }
  27438. batch_rule(self, out);
  27439. }
  27440. template <typename batch_rule_t, batch_rule_t batch_rule>
  27441. void _foreach_sign_out_generated_plumbing(at::TensorList self, at::TensorList out) {
  27442. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27443. auto maybe_layer = maybeCurrentDynamicLayer();
  27444. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27445. int64_t cur_level = maybe_layer->layerId();
  27446. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27447. return at::_ops::_foreach_sign_out::call(self, out);
  27448. }
  27449. batch_rule(self, out);
  27450. }
  27451. template <typename batch_rule_t, batch_rule_t batch_rule>
  27452. void _foreach_sin_out_generated_plumbing(at::TensorList self, at::TensorList out) {
  27453. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27454. auto maybe_layer = maybeCurrentDynamicLayer();
  27455. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27456. int64_t cur_level = maybe_layer->layerId();
  27457. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27458. return at::_ops::_foreach_sin_out::call(self, out);
  27459. }
  27460. batch_rule(self, out);
  27461. }
  27462. template <typename batch_rule_t, batch_rule_t batch_rule>
  27463. void _foreach_sinh_out_generated_plumbing(at::TensorList self, at::TensorList out) {
  27464. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27465. auto maybe_layer = maybeCurrentDynamicLayer();
  27466. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27467. int64_t cur_level = maybe_layer->layerId();
  27468. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27469. return at::_ops::_foreach_sinh_out::call(self, out);
  27470. }
  27471. batch_rule(self, out);
  27472. }
  27473. template <typename batch_rule_t, batch_rule_t batch_rule>
  27474. void _foreach_sqrt_out_generated_plumbing(at::TensorList self, at::TensorList out) {
  27475. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27476. auto maybe_layer = maybeCurrentDynamicLayer();
  27477. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27478. int64_t cur_level = maybe_layer->layerId();
  27479. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27480. return at::_ops::_foreach_sqrt_out::call(self, out);
  27481. }
  27482. batch_rule(self, out);
  27483. }
  27484. template <typename batch_rule_t, batch_rule_t batch_rule>
  27485. void _foreach_tan_out_generated_plumbing(at::TensorList self, at::TensorList out) {
  27486. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27487. auto maybe_layer = maybeCurrentDynamicLayer();
  27488. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27489. int64_t cur_level = maybe_layer->layerId();
  27490. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27491. return at::_ops::_foreach_tan_out::call(self, out);
  27492. }
  27493. batch_rule(self, out);
  27494. }
  27495. template <typename batch_rule_t, batch_rule_t batch_rule>
  27496. void _foreach_tanh_out_generated_plumbing(at::TensorList self, at::TensorList out) {
  27497. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27498. auto maybe_layer = maybeCurrentDynamicLayer();
  27499. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27500. int64_t cur_level = maybe_layer->layerId();
  27501. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27502. return at::_ops::_foreach_tanh_out::call(self, out);
  27503. }
  27504. batch_rule(self, out);
  27505. }
  27506. template <typename batch_rule_t, batch_rule_t batch_rule>
  27507. void _foreach_trunc_out_generated_plumbing(at::TensorList self, at::TensorList out) {
  27508. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27509. auto maybe_layer = maybeCurrentDynamicLayer();
  27510. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27511. int64_t cur_level = maybe_layer->layerId();
  27512. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27513. return at::_ops::_foreach_trunc_out::call(self, out);
  27514. }
  27515. batch_rule(self, out);
  27516. }
  27517. template <typename batch_rule_t, batch_rule_t batch_rule>
  27518. void _foreach_zero_out_generated_plumbing(at::TensorList self, at::TensorList out) {
  27519. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27520. auto maybe_layer = maybeCurrentDynamicLayer();
  27521. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27522. int64_t cur_level = maybe_layer->layerId();
  27523. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27524. return at::_ops::_foreach_zero_out::call(self, out);
  27525. }
  27526. batch_rule(self, out);
  27527. }
  27528. template <typename batch_rule_t, batch_rule_t batch_rule>
  27529. ::std::vector<at::Tensor> _foreach_zero_generated_plumbing(at::TensorList self) {
  27530. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27531. auto maybe_layer = maybeCurrentDynamicLayer();
  27532. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  27533. int64_t cur_level = maybe_layer->layerId();
  27534. if (!isBatchedAtLevel(self, cur_level)) {
  27535. return at::_ops::_foreach_zero::call(self);
  27536. }
  27537. auto results = batch_rule(self);
  27538. return makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level);
  27539. }
  27540. template <typename batch_rule_t, batch_rule_t batch_rule>
  27541. void _foreach_copy_out_generated_plumbing(at::TensorList self, at::TensorList src, bool non_blocking, at::TensorList out) {
  27542. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27543. auto maybe_layer = maybeCurrentDynamicLayer();
  27544. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27545. int64_t cur_level = maybe_layer->layerId();
  27546. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(src, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27547. return at::_ops::_foreach_copy_out::call(self, src, non_blocking, out);
  27548. }
  27549. batch_rule(self, src, non_blocking, out);
  27550. }
  27551. template <typename batch_rule_t, batch_rule_t batch_rule>
  27552. ::std::tuple<at::Tensor,at::Tensor> rrelu_with_noise_functional_generated_plumbing(const at::Tensor & self, const at::Tensor & noise, const at::Scalar & lower, const at::Scalar & upper, bool training, ::std::optional<at::Generator> generator) {
  27553. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27554. auto maybe_layer = maybeCurrentDynamicLayer();
  27555. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  27556. int64_t cur_level = maybe_layer->layerId();
  27557. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(noise, cur_level)) {
  27558. return at::_ops::rrelu_with_noise_functional::call(self, noise, lower, upper, training, generator);
  27559. }
  27560. auto [self_value, self_bdim] = unwrapTensorAtLevel(self, cur_level);
  27561. auto [noise_value, noise_bdim] = unwrapTensorAtLevel(noise, cur_level);
  27562. auto results = batch_rule(self_value, self_bdim, noise_value, noise_bdim, lower, upper, training, generator);
  27563. return std::make_tuple(makeBatched(std::get<0>(results), std::get<1>(results), cur_level), makeBatched(std::get<2>(results), std::get<3>(results), cur_level));
  27564. }
  27565. template <typename batch_rule_t, batch_rule_t batch_rule>
  27566. void _fused_adam_out_generated_plumbing(at::TensorList self, at::TensorList grads, at::TensorList exp_avgs, at::TensorList exp_avg_sqs, at::TensorList max_exp_avg_sqs, at::TensorList state_steps, double lr, double beta1, double beta2, double weight_decay, double eps, bool amsgrad, bool maximize, const ::std::optional<at::Tensor> & grad_scale, const ::std::optional<at::Tensor> & found_inf, at::TensorList out) {
  27567. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27568. auto maybe_layer = maybeCurrentDynamicLayer();
  27569. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27570. int64_t cur_level = maybe_layer->layerId();
  27571. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(grads, cur_level) && !isBatchedAtLevel(exp_avgs, cur_level) && !isBatchedAtLevel(exp_avg_sqs, cur_level) && !isBatchedAtLevel(max_exp_avg_sqs, cur_level) && !isBatchedAtLevel(state_steps, cur_level) && !isBatchedAtLevel(grad_scale, cur_level) && !isBatchedAtLevel(found_inf, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27572. return at::_ops::_fused_adam_out::call(self, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps, lr, beta1, beta2, weight_decay, eps, amsgrad, maximize, grad_scale, found_inf, out);
  27573. }
  27574. std::optional<Tensor> grad_scale_value;
  27575. std::optional<int64_t> grad_scale_bdim;
  27576. if (grad_scale) {
  27577. std::tie(grad_scale_value, grad_scale_bdim) = unwrapTensorAtLevel(grad_scale.value(), cur_level);
  27578. }
  27579. std::optional<Tensor> found_inf_value;
  27580. std::optional<int64_t> found_inf_bdim;
  27581. if (found_inf) {
  27582. std::tie(found_inf_value, found_inf_bdim) = unwrapTensorAtLevel(found_inf.value(), cur_level);
  27583. }
  27584. batch_rule(self, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps, lr, beta1, beta2, weight_decay, eps, amsgrad, maximize, grad_scale_value, grad_scale_bdim, found_inf_value, found_inf_bdim, out);
  27585. }
  27586. template <typename batch_rule_t, batch_rule_t batch_rule>
  27587. ::std::tuple<::std::vector<at::Tensor>,::std::vector<at::Tensor>,::std::vector<at::Tensor>,::std::vector<at::Tensor>,::std::vector<at::Tensor>> _fused_adam_generated_plumbing(at::TensorList self, at::TensorList grads, at::TensorList exp_avgs, at::TensorList exp_avg_sqs, at::TensorList max_exp_avg_sqs, at::TensorList state_steps, double lr, double beta1, double beta2, double weight_decay, double eps, bool amsgrad, bool maximize, const ::std::optional<at::Tensor> & grad_scale, const ::std::optional<at::Tensor> & found_inf) {
  27588. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27589. auto maybe_layer = maybeCurrentDynamicLayer();
  27590. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  27591. int64_t cur_level = maybe_layer->layerId();
  27592. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(grads, cur_level) && !isBatchedAtLevel(exp_avgs, cur_level) && !isBatchedAtLevel(exp_avg_sqs, cur_level) && !isBatchedAtLevel(max_exp_avg_sqs, cur_level) && !isBatchedAtLevel(state_steps, cur_level) && !isBatchedAtLevel(grad_scale, cur_level) && !isBatchedAtLevel(found_inf, cur_level)) {
  27593. return at::_ops::_fused_adam::call(self, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps, lr, beta1, beta2, weight_decay, eps, amsgrad, maximize, grad_scale, found_inf);
  27594. }
  27595. std::optional<Tensor> grad_scale_value;
  27596. std::optional<int64_t> grad_scale_bdim;
  27597. if (grad_scale) {
  27598. std::tie(grad_scale_value, grad_scale_bdim) = unwrapTensorAtLevel(grad_scale.value(), cur_level);
  27599. }
  27600. std::optional<Tensor> found_inf_value;
  27601. std::optional<int64_t> found_inf_bdim;
  27602. if (found_inf) {
  27603. std::tie(found_inf_value, found_inf_bdim) = unwrapTensorAtLevel(found_inf.value(), cur_level);
  27604. }
  27605. auto results = batch_rule(self, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps, lr, beta1, beta2, weight_decay, eps, amsgrad, maximize, grad_scale_value, grad_scale_bdim, found_inf_value, found_inf_bdim);
  27606. return std::make_tuple(makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level), makeBatchedVector(std::get<2>(results), std::get<3>(results), cur_level), makeBatchedVector(std::get<4>(results), std::get<5>(results), cur_level), makeBatchedVector(std::get<6>(results), std::get<7>(results), cur_level), makeBatchedVector(std::get<8>(results), std::get<9>(results), cur_level));
  27607. }
  27608. template <typename batch_rule_t, batch_rule_t batch_rule>
  27609. void _fused_adam_tensor_lr_out_generated_plumbing(at::TensorList self, at::TensorList grads, at::TensorList exp_avgs, at::TensorList exp_avg_sqs, at::TensorList max_exp_avg_sqs, at::TensorList state_steps, const at::Tensor & lr, double beta1, double beta2, double weight_decay, double eps, bool amsgrad, bool maximize, const ::std::optional<at::Tensor> & grad_scale, const ::std::optional<at::Tensor> & found_inf, at::TensorList out) {
  27610. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27611. auto maybe_layer = maybeCurrentDynamicLayer();
  27612. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27613. int64_t cur_level = maybe_layer->layerId();
  27614. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(grads, cur_level) && !isBatchedAtLevel(exp_avgs, cur_level) && !isBatchedAtLevel(exp_avg_sqs, cur_level) && !isBatchedAtLevel(max_exp_avg_sqs, cur_level) && !isBatchedAtLevel(state_steps, cur_level) && !isBatchedAtLevel(lr, cur_level) && !isBatchedAtLevel(grad_scale, cur_level) && !isBatchedAtLevel(found_inf, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27615. return at::_ops::_fused_adam_tensor_lr_out::call(self, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps, lr, beta1, beta2, weight_decay, eps, amsgrad, maximize, grad_scale, found_inf, out);
  27616. }
  27617. auto [lr_value, lr_bdim] = unwrapTensorAtLevel(lr, cur_level);
  27618. std::optional<Tensor> grad_scale_value;
  27619. std::optional<int64_t> grad_scale_bdim;
  27620. if (grad_scale) {
  27621. std::tie(grad_scale_value, grad_scale_bdim) = unwrapTensorAtLevel(grad_scale.value(), cur_level);
  27622. }
  27623. std::optional<Tensor> found_inf_value;
  27624. std::optional<int64_t> found_inf_bdim;
  27625. if (found_inf) {
  27626. std::tie(found_inf_value, found_inf_bdim) = unwrapTensorAtLevel(found_inf.value(), cur_level);
  27627. }
  27628. batch_rule(self, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps, lr_value, lr_bdim, beta1, beta2, weight_decay, eps, amsgrad, maximize, grad_scale_value, grad_scale_bdim, found_inf_value, found_inf_bdim, out);
  27629. }
  27630. template <typename batch_rule_t, batch_rule_t batch_rule>
  27631. ::std::tuple<::std::vector<at::Tensor>,::std::vector<at::Tensor>,::std::vector<at::Tensor>,::std::vector<at::Tensor>,::std::vector<at::Tensor>> _fused_adam_tensor_lr_generated_plumbing(at::TensorList self, at::TensorList grads, at::TensorList exp_avgs, at::TensorList exp_avg_sqs, at::TensorList max_exp_avg_sqs, at::TensorList state_steps, const at::Tensor & lr, double beta1, double beta2, double weight_decay, double eps, bool amsgrad, bool maximize, const ::std::optional<at::Tensor> & grad_scale, const ::std::optional<at::Tensor> & found_inf) {
  27632. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27633. auto maybe_layer = maybeCurrentDynamicLayer();
  27634. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  27635. int64_t cur_level = maybe_layer->layerId();
  27636. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(grads, cur_level) && !isBatchedAtLevel(exp_avgs, cur_level) && !isBatchedAtLevel(exp_avg_sqs, cur_level) && !isBatchedAtLevel(max_exp_avg_sqs, cur_level) && !isBatchedAtLevel(state_steps, cur_level) && !isBatchedAtLevel(lr, cur_level) && !isBatchedAtLevel(grad_scale, cur_level) && !isBatchedAtLevel(found_inf, cur_level)) {
  27637. return at::_ops::_fused_adam_tensor_lr::call(self, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps, lr, beta1, beta2, weight_decay, eps, amsgrad, maximize, grad_scale, found_inf);
  27638. }
  27639. auto [lr_value, lr_bdim] = unwrapTensorAtLevel(lr, cur_level);
  27640. std::optional<Tensor> grad_scale_value;
  27641. std::optional<int64_t> grad_scale_bdim;
  27642. if (grad_scale) {
  27643. std::tie(grad_scale_value, grad_scale_bdim) = unwrapTensorAtLevel(grad_scale.value(), cur_level);
  27644. }
  27645. std::optional<Tensor> found_inf_value;
  27646. std::optional<int64_t> found_inf_bdim;
  27647. if (found_inf) {
  27648. std::tie(found_inf_value, found_inf_bdim) = unwrapTensorAtLevel(found_inf.value(), cur_level);
  27649. }
  27650. auto results = batch_rule(self, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps, lr_value, lr_bdim, beta1, beta2, weight_decay, eps, amsgrad, maximize, grad_scale_value, grad_scale_bdim, found_inf_value, found_inf_bdim);
  27651. return std::make_tuple(makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level), makeBatchedVector(std::get<2>(results), std::get<3>(results), cur_level), makeBatchedVector(std::get<4>(results), std::get<5>(results), cur_level), makeBatchedVector(std::get<6>(results), std::get<7>(results), cur_level), makeBatchedVector(std::get<8>(results), std::get<9>(results), cur_level));
  27652. }
  27653. template <typename batch_rule_t, batch_rule_t batch_rule>
  27654. void _fused_adamw_out_generated_plumbing(at::TensorList self, at::TensorList grads, at::TensorList exp_avgs, at::TensorList exp_avg_sqs, at::TensorList max_exp_avg_sqs, at::TensorList state_steps, double lr, double beta1, double beta2, double weight_decay, double eps, bool amsgrad, bool maximize, const ::std::optional<at::Tensor> & grad_scale, const ::std::optional<at::Tensor> & found_inf, at::TensorList out) {
  27655. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27656. auto maybe_layer = maybeCurrentDynamicLayer();
  27657. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27658. int64_t cur_level = maybe_layer->layerId();
  27659. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(grads, cur_level) && !isBatchedAtLevel(exp_avgs, cur_level) && !isBatchedAtLevel(exp_avg_sqs, cur_level) && !isBatchedAtLevel(max_exp_avg_sqs, cur_level) && !isBatchedAtLevel(state_steps, cur_level) && !isBatchedAtLevel(grad_scale, cur_level) && !isBatchedAtLevel(found_inf, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27660. return at::_ops::_fused_adamw_out::call(self, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps, lr, beta1, beta2, weight_decay, eps, amsgrad, maximize, grad_scale, found_inf, out);
  27661. }
  27662. std::optional<Tensor> grad_scale_value;
  27663. std::optional<int64_t> grad_scale_bdim;
  27664. if (grad_scale) {
  27665. std::tie(grad_scale_value, grad_scale_bdim) = unwrapTensorAtLevel(grad_scale.value(), cur_level);
  27666. }
  27667. std::optional<Tensor> found_inf_value;
  27668. std::optional<int64_t> found_inf_bdim;
  27669. if (found_inf) {
  27670. std::tie(found_inf_value, found_inf_bdim) = unwrapTensorAtLevel(found_inf.value(), cur_level);
  27671. }
  27672. batch_rule(self, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps, lr, beta1, beta2, weight_decay, eps, amsgrad, maximize, grad_scale_value, grad_scale_bdim, found_inf_value, found_inf_bdim, out);
  27673. }
  27674. template <typename batch_rule_t, batch_rule_t batch_rule>
  27675. ::std::tuple<::std::vector<at::Tensor>,::std::vector<at::Tensor>,::std::vector<at::Tensor>,::std::vector<at::Tensor>,::std::vector<at::Tensor>> _fused_adamw_generated_plumbing(at::TensorList self, at::TensorList grads, at::TensorList exp_avgs, at::TensorList exp_avg_sqs, at::TensorList max_exp_avg_sqs, at::TensorList state_steps, double lr, double beta1, double beta2, double weight_decay, double eps, bool amsgrad, bool maximize, const ::std::optional<at::Tensor> & grad_scale, const ::std::optional<at::Tensor> & found_inf) {
  27676. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27677. auto maybe_layer = maybeCurrentDynamicLayer();
  27678. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  27679. int64_t cur_level = maybe_layer->layerId();
  27680. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(grads, cur_level) && !isBatchedAtLevel(exp_avgs, cur_level) && !isBatchedAtLevel(exp_avg_sqs, cur_level) && !isBatchedAtLevel(max_exp_avg_sqs, cur_level) && !isBatchedAtLevel(state_steps, cur_level) && !isBatchedAtLevel(grad_scale, cur_level) && !isBatchedAtLevel(found_inf, cur_level)) {
  27681. return at::_ops::_fused_adamw::call(self, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps, lr, beta1, beta2, weight_decay, eps, amsgrad, maximize, grad_scale, found_inf);
  27682. }
  27683. std::optional<Tensor> grad_scale_value;
  27684. std::optional<int64_t> grad_scale_bdim;
  27685. if (grad_scale) {
  27686. std::tie(grad_scale_value, grad_scale_bdim) = unwrapTensorAtLevel(grad_scale.value(), cur_level);
  27687. }
  27688. std::optional<Tensor> found_inf_value;
  27689. std::optional<int64_t> found_inf_bdim;
  27690. if (found_inf) {
  27691. std::tie(found_inf_value, found_inf_bdim) = unwrapTensorAtLevel(found_inf.value(), cur_level);
  27692. }
  27693. auto results = batch_rule(self, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps, lr, beta1, beta2, weight_decay, eps, amsgrad, maximize, grad_scale_value, grad_scale_bdim, found_inf_value, found_inf_bdim);
  27694. return std::make_tuple(makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level), makeBatchedVector(std::get<2>(results), std::get<3>(results), cur_level), makeBatchedVector(std::get<4>(results), std::get<5>(results), cur_level), makeBatchedVector(std::get<6>(results), std::get<7>(results), cur_level), makeBatchedVector(std::get<8>(results), std::get<9>(results), cur_level));
  27695. }
  27696. template <typename batch_rule_t, batch_rule_t batch_rule>
  27697. void _fused_adamw_tensor_lr_out_generated_plumbing(at::TensorList self, at::TensorList grads, at::TensorList exp_avgs, at::TensorList exp_avg_sqs, at::TensorList max_exp_avg_sqs, at::TensorList state_steps, const at::Tensor & lr, double beta1, double beta2, double weight_decay, double eps, bool amsgrad, bool maximize, const ::std::optional<at::Tensor> & grad_scale, const ::std::optional<at::Tensor> & found_inf, at::TensorList out) {
  27698. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27699. auto maybe_layer = maybeCurrentDynamicLayer();
  27700. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27701. int64_t cur_level = maybe_layer->layerId();
  27702. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(grads, cur_level) && !isBatchedAtLevel(exp_avgs, cur_level) && !isBatchedAtLevel(exp_avg_sqs, cur_level) && !isBatchedAtLevel(max_exp_avg_sqs, cur_level) && !isBatchedAtLevel(state_steps, cur_level) && !isBatchedAtLevel(lr, cur_level) && !isBatchedAtLevel(grad_scale, cur_level) && !isBatchedAtLevel(found_inf, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27703. return at::_ops::_fused_adamw_tensor_lr_out::call(self, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps, lr, beta1, beta2, weight_decay, eps, amsgrad, maximize, grad_scale, found_inf, out);
  27704. }
  27705. auto [lr_value, lr_bdim] = unwrapTensorAtLevel(lr, cur_level);
  27706. std::optional<Tensor> grad_scale_value;
  27707. std::optional<int64_t> grad_scale_bdim;
  27708. if (grad_scale) {
  27709. std::tie(grad_scale_value, grad_scale_bdim) = unwrapTensorAtLevel(grad_scale.value(), cur_level);
  27710. }
  27711. std::optional<Tensor> found_inf_value;
  27712. std::optional<int64_t> found_inf_bdim;
  27713. if (found_inf) {
  27714. std::tie(found_inf_value, found_inf_bdim) = unwrapTensorAtLevel(found_inf.value(), cur_level);
  27715. }
  27716. batch_rule(self, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps, lr_value, lr_bdim, beta1, beta2, weight_decay, eps, amsgrad, maximize, grad_scale_value, grad_scale_bdim, found_inf_value, found_inf_bdim, out);
  27717. }
  27718. template <typename batch_rule_t, batch_rule_t batch_rule>
  27719. ::std::tuple<::std::vector<at::Tensor>,::std::vector<at::Tensor>,::std::vector<at::Tensor>,::std::vector<at::Tensor>,::std::vector<at::Tensor>> _fused_adamw_tensor_lr_generated_plumbing(at::TensorList self, at::TensorList grads, at::TensorList exp_avgs, at::TensorList exp_avg_sqs, at::TensorList max_exp_avg_sqs, at::TensorList state_steps, const at::Tensor & lr, double beta1, double beta2, double weight_decay, double eps, bool amsgrad, bool maximize, const ::std::optional<at::Tensor> & grad_scale, const ::std::optional<at::Tensor> & found_inf) {
  27720. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27721. auto maybe_layer = maybeCurrentDynamicLayer();
  27722. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  27723. int64_t cur_level = maybe_layer->layerId();
  27724. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(grads, cur_level) && !isBatchedAtLevel(exp_avgs, cur_level) && !isBatchedAtLevel(exp_avg_sqs, cur_level) && !isBatchedAtLevel(max_exp_avg_sqs, cur_level) && !isBatchedAtLevel(state_steps, cur_level) && !isBatchedAtLevel(lr, cur_level) && !isBatchedAtLevel(grad_scale, cur_level) && !isBatchedAtLevel(found_inf, cur_level)) {
  27725. return at::_ops::_fused_adamw_tensor_lr::call(self, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps, lr, beta1, beta2, weight_decay, eps, amsgrad, maximize, grad_scale, found_inf);
  27726. }
  27727. auto [lr_value, lr_bdim] = unwrapTensorAtLevel(lr, cur_level);
  27728. std::optional<Tensor> grad_scale_value;
  27729. std::optional<int64_t> grad_scale_bdim;
  27730. if (grad_scale) {
  27731. std::tie(grad_scale_value, grad_scale_bdim) = unwrapTensorAtLevel(grad_scale.value(), cur_level);
  27732. }
  27733. std::optional<Tensor> found_inf_value;
  27734. std::optional<int64_t> found_inf_bdim;
  27735. if (found_inf) {
  27736. std::tie(found_inf_value, found_inf_bdim) = unwrapTensorAtLevel(found_inf.value(), cur_level);
  27737. }
  27738. auto results = batch_rule(self, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps, lr_value, lr_bdim, beta1, beta2, weight_decay, eps, amsgrad, maximize, grad_scale_value, grad_scale_bdim, found_inf_value, found_inf_bdim);
  27739. return std::make_tuple(makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level), makeBatchedVector(std::get<2>(results), std::get<3>(results), cur_level), makeBatchedVector(std::get<4>(results), std::get<5>(results), cur_level), makeBatchedVector(std::get<6>(results), std::get<7>(results), cur_level), makeBatchedVector(std::get<8>(results), std::get<9>(results), cur_level));
  27740. }
  27741. template <typename batch_rule_t, batch_rule_t batch_rule>
  27742. void _fused_sgd_out_generated_plumbing(at::TensorList self, at::TensorList grads, at::TensorList momentum_buffer_list, double weight_decay, double momentum, double lr, double dampening, bool nesterov, bool maximize, bool is_first_step, const ::std::optional<at::Tensor> & grad_scale, const ::std::optional<at::Tensor> & found_inf, at::TensorList out) {
  27743. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27744. auto maybe_layer = maybeCurrentDynamicLayer();
  27745. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27746. int64_t cur_level = maybe_layer->layerId();
  27747. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(grads, cur_level) && !isBatchedAtLevel(momentum_buffer_list, cur_level) && !isBatchedAtLevel(grad_scale, cur_level) && !isBatchedAtLevel(found_inf, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27748. return at::_ops::_fused_sgd_out::call(self, grads, momentum_buffer_list, weight_decay, momentum, lr, dampening, nesterov, maximize, is_first_step, grad_scale, found_inf, out);
  27749. }
  27750. std::optional<Tensor> grad_scale_value;
  27751. std::optional<int64_t> grad_scale_bdim;
  27752. if (grad_scale) {
  27753. std::tie(grad_scale_value, grad_scale_bdim) = unwrapTensorAtLevel(grad_scale.value(), cur_level);
  27754. }
  27755. std::optional<Tensor> found_inf_value;
  27756. std::optional<int64_t> found_inf_bdim;
  27757. if (found_inf) {
  27758. std::tie(found_inf_value, found_inf_bdim) = unwrapTensorAtLevel(found_inf.value(), cur_level);
  27759. }
  27760. batch_rule(self, grads, momentum_buffer_list, weight_decay, momentum, lr, dampening, nesterov, maximize, is_first_step, grad_scale_value, grad_scale_bdim, found_inf_value, found_inf_bdim, out);
  27761. }
  27762. template <typename batch_rule_t, batch_rule_t batch_rule>
  27763. ::std::tuple<::std::vector<at::Tensor>,::std::vector<at::Tensor>,::std::vector<at::Tensor>> _fused_sgd_generated_plumbing(at::TensorList self, at::TensorList grads, at::TensorList momentum_buffer_list, double weight_decay, double momentum, double lr, double dampening, bool nesterov, bool maximize, bool is_first_step, const ::std::optional<at::Tensor> & grad_scale, const ::std::optional<at::Tensor> & found_inf) {
  27764. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27765. auto maybe_layer = maybeCurrentDynamicLayer();
  27766. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  27767. int64_t cur_level = maybe_layer->layerId();
  27768. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(grads, cur_level) && !isBatchedAtLevel(momentum_buffer_list, cur_level) && !isBatchedAtLevel(grad_scale, cur_level) && !isBatchedAtLevel(found_inf, cur_level)) {
  27769. return at::_ops::_fused_sgd::call(self, grads, momentum_buffer_list, weight_decay, momentum, lr, dampening, nesterov, maximize, is_first_step, grad_scale, found_inf);
  27770. }
  27771. std::optional<Tensor> grad_scale_value;
  27772. std::optional<int64_t> grad_scale_bdim;
  27773. if (grad_scale) {
  27774. std::tie(grad_scale_value, grad_scale_bdim) = unwrapTensorAtLevel(grad_scale.value(), cur_level);
  27775. }
  27776. std::optional<Tensor> found_inf_value;
  27777. std::optional<int64_t> found_inf_bdim;
  27778. if (found_inf) {
  27779. std::tie(found_inf_value, found_inf_bdim) = unwrapTensorAtLevel(found_inf.value(), cur_level);
  27780. }
  27781. auto results = batch_rule(self, grads, momentum_buffer_list, weight_decay, momentum, lr, dampening, nesterov, maximize, is_first_step, grad_scale_value, grad_scale_bdim, found_inf_value, found_inf_bdim);
  27782. return std::make_tuple(makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level), makeBatchedVector(std::get<2>(results), std::get<3>(results), cur_level), makeBatchedVector(std::get<4>(results), std::get<5>(results), cur_level));
  27783. }
  27784. template <typename batch_rule_t, batch_rule_t batch_rule>
  27785. void _fused_sgd_tensor_lr_out_generated_plumbing(at::TensorList self, at::TensorList grads, at::TensorList momentum_buffer_list, double weight_decay, double momentum, const at::Tensor & lr, double dampening, bool nesterov, bool maximize, bool is_first_step, const ::std::optional<at::Tensor> & grad_scale, const ::std::optional<at::Tensor> & found_inf, at::TensorList out) {
  27786. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27787. auto maybe_layer = maybeCurrentDynamicLayer();
  27788. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27789. int64_t cur_level = maybe_layer->layerId();
  27790. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(grads, cur_level) && !isBatchedAtLevel(momentum_buffer_list, cur_level) && !isBatchedAtLevel(lr, cur_level) && !isBatchedAtLevel(grad_scale, cur_level) && !isBatchedAtLevel(found_inf, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27791. return at::_ops::_fused_sgd_tensor_lr_out::call(self, grads, momentum_buffer_list, weight_decay, momentum, lr, dampening, nesterov, maximize, is_first_step, grad_scale, found_inf, out);
  27792. }
  27793. auto [lr_value, lr_bdim] = unwrapTensorAtLevel(lr, cur_level);
  27794. std::optional<Tensor> grad_scale_value;
  27795. std::optional<int64_t> grad_scale_bdim;
  27796. if (grad_scale) {
  27797. std::tie(grad_scale_value, grad_scale_bdim) = unwrapTensorAtLevel(grad_scale.value(), cur_level);
  27798. }
  27799. std::optional<Tensor> found_inf_value;
  27800. std::optional<int64_t> found_inf_bdim;
  27801. if (found_inf) {
  27802. std::tie(found_inf_value, found_inf_bdim) = unwrapTensorAtLevel(found_inf.value(), cur_level);
  27803. }
  27804. batch_rule(self, grads, momentum_buffer_list, weight_decay, momentum, lr_value, lr_bdim, dampening, nesterov, maximize, is_first_step, grad_scale_value, grad_scale_bdim, found_inf_value, found_inf_bdim, out);
  27805. }
  27806. template <typename batch_rule_t, batch_rule_t batch_rule>
  27807. ::std::tuple<::std::vector<at::Tensor>,::std::vector<at::Tensor>,::std::vector<at::Tensor>> _fused_sgd_tensor_lr_generated_plumbing(at::TensorList self, at::TensorList grads, at::TensorList momentum_buffer_list, double weight_decay, double momentum, const at::Tensor & lr, double dampening, bool nesterov, bool maximize, bool is_first_step, const ::std::optional<at::Tensor> & grad_scale, const ::std::optional<at::Tensor> & found_inf) {
  27808. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27809. auto maybe_layer = maybeCurrentDynamicLayer();
  27810. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  27811. int64_t cur_level = maybe_layer->layerId();
  27812. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(grads, cur_level) && !isBatchedAtLevel(momentum_buffer_list, cur_level) && !isBatchedAtLevel(lr, cur_level) && !isBatchedAtLevel(grad_scale, cur_level) && !isBatchedAtLevel(found_inf, cur_level)) {
  27813. return at::_ops::_fused_sgd_tensor_lr::call(self, grads, momentum_buffer_list, weight_decay, momentum, lr, dampening, nesterov, maximize, is_first_step, grad_scale, found_inf);
  27814. }
  27815. auto [lr_value, lr_bdim] = unwrapTensorAtLevel(lr, cur_level);
  27816. std::optional<Tensor> grad_scale_value;
  27817. std::optional<int64_t> grad_scale_bdim;
  27818. if (grad_scale) {
  27819. std::tie(grad_scale_value, grad_scale_bdim) = unwrapTensorAtLevel(grad_scale.value(), cur_level);
  27820. }
  27821. std::optional<Tensor> found_inf_value;
  27822. std::optional<int64_t> found_inf_bdim;
  27823. if (found_inf) {
  27824. std::tie(found_inf_value, found_inf_bdim) = unwrapTensorAtLevel(found_inf.value(), cur_level);
  27825. }
  27826. auto results = batch_rule(self, grads, momentum_buffer_list, weight_decay, momentum, lr_value, lr_bdim, dampening, nesterov, maximize, is_first_step, grad_scale_value, grad_scale_bdim, found_inf_value, found_inf_bdim);
  27827. return std::make_tuple(makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level), makeBatchedVector(std::get<2>(results), std::get<3>(results), cur_level), makeBatchedVector(std::get<4>(results), std::get<5>(results), cur_level));
  27828. }
  27829. template <typename batch_rule_t, batch_rule_t batch_rule>
  27830. void _fused_adagrad_out_generated_plumbing(at::TensorList self, at::TensorList grads, at::TensorList state_sums, at::TensorList state_steps, double lr, double lr_decay, double weight_decay, double eps, bool maximize, const ::std::optional<at::Tensor> & grad_scale, const ::std::optional<at::Tensor> & found_inf, at::TensorList out) {
  27831. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27832. auto maybe_layer = maybeCurrentDynamicLayer();
  27833. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27834. int64_t cur_level = maybe_layer->layerId();
  27835. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(grads, cur_level) && !isBatchedAtLevel(state_sums, cur_level) && !isBatchedAtLevel(state_steps, cur_level) && !isBatchedAtLevel(grad_scale, cur_level) && !isBatchedAtLevel(found_inf, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27836. return at::_ops::_fused_adagrad_out::call(self, grads, state_sums, state_steps, lr, lr_decay, weight_decay, eps, maximize, grad_scale, found_inf, out);
  27837. }
  27838. std::optional<Tensor> grad_scale_value;
  27839. std::optional<int64_t> grad_scale_bdim;
  27840. if (grad_scale) {
  27841. std::tie(grad_scale_value, grad_scale_bdim) = unwrapTensorAtLevel(grad_scale.value(), cur_level);
  27842. }
  27843. std::optional<Tensor> found_inf_value;
  27844. std::optional<int64_t> found_inf_bdim;
  27845. if (found_inf) {
  27846. std::tie(found_inf_value, found_inf_bdim) = unwrapTensorAtLevel(found_inf.value(), cur_level);
  27847. }
  27848. batch_rule(self, grads, state_sums, state_steps, lr, lr_decay, weight_decay, eps, maximize, grad_scale_value, grad_scale_bdim, found_inf_value, found_inf_bdim, out);
  27849. }
  27850. template <typename batch_rule_t, batch_rule_t batch_rule>
  27851. ::std::tuple<::std::vector<at::Tensor>,::std::vector<at::Tensor>,::std::vector<at::Tensor>,::std::vector<at::Tensor>> _fused_adagrad_generated_plumbing(at::TensorList self, at::TensorList grads, at::TensorList state_sums, at::TensorList state_steps, double lr, double lr_decay, double weight_decay, double eps, bool maximize, const ::std::optional<at::Tensor> & grad_scale, const ::std::optional<at::Tensor> & found_inf) {
  27852. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27853. auto maybe_layer = maybeCurrentDynamicLayer();
  27854. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  27855. int64_t cur_level = maybe_layer->layerId();
  27856. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(grads, cur_level) && !isBatchedAtLevel(state_sums, cur_level) && !isBatchedAtLevel(state_steps, cur_level) && !isBatchedAtLevel(grad_scale, cur_level) && !isBatchedAtLevel(found_inf, cur_level)) {
  27857. return at::_ops::_fused_adagrad::call(self, grads, state_sums, state_steps, lr, lr_decay, weight_decay, eps, maximize, grad_scale, found_inf);
  27858. }
  27859. std::optional<Tensor> grad_scale_value;
  27860. std::optional<int64_t> grad_scale_bdim;
  27861. if (grad_scale) {
  27862. std::tie(grad_scale_value, grad_scale_bdim) = unwrapTensorAtLevel(grad_scale.value(), cur_level);
  27863. }
  27864. std::optional<Tensor> found_inf_value;
  27865. std::optional<int64_t> found_inf_bdim;
  27866. if (found_inf) {
  27867. std::tie(found_inf_value, found_inf_bdim) = unwrapTensorAtLevel(found_inf.value(), cur_level);
  27868. }
  27869. auto results = batch_rule(self, grads, state_sums, state_steps, lr, lr_decay, weight_decay, eps, maximize, grad_scale_value, grad_scale_bdim, found_inf_value, found_inf_bdim);
  27870. return std::make_tuple(makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level), makeBatchedVector(std::get<2>(results), std::get<3>(results), cur_level), makeBatchedVector(std::get<4>(results), std::get<5>(results), cur_level), makeBatchedVector(std::get<6>(results), std::get<7>(results), cur_level));
  27871. }
  27872. template <typename batch_rule_t, batch_rule_t batch_rule>
  27873. void _fused_adagrad_tensor_lr_out_generated_plumbing(at::TensorList self, at::TensorList grads, at::TensorList state_sums, at::TensorList state_steps, const at::Tensor & lr, double lr_decay, double weight_decay, double eps, bool maximize, const ::std::optional<at::Tensor> & grad_scale, const ::std::optional<at::Tensor> & found_inf, at::TensorList out) {
  27874. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27875. auto maybe_layer = maybeCurrentDynamicLayer();
  27876. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns");
  27877. int64_t cur_level = maybe_layer->layerId();
  27878. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(grads, cur_level) && !isBatchedAtLevel(state_sums, cur_level) && !isBatchedAtLevel(state_steps, cur_level) && !isBatchedAtLevel(lr, cur_level) && !isBatchedAtLevel(grad_scale, cur_level) && !isBatchedAtLevel(found_inf, cur_level) && !isBatchedAtLevel(out, cur_level)) {
  27879. return at::_ops::_fused_adagrad_tensor_lr_out::call(self, grads, state_sums, state_steps, lr, lr_decay, weight_decay, eps, maximize, grad_scale, found_inf, out);
  27880. }
  27881. auto [lr_value, lr_bdim] = unwrapTensorAtLevel(lr, cur_level);
  27882. std::optional<Tensor> grad_scale_value;
  27883. std::optional<int64_t> grad_scale_bdim;
  27884. if (grad_scale) {
  27885. std::tie(grad_scale_value, grad_scale_bdim) = unwrapTensorAtLevel(grad_scale.value(), cur_level);
  27886. }
  27887. std::optional<Tensor> found_inf_value;
  27888. std::optional<int64_t> found_inf_bdim;
  27889. if (found_inf) {
  27890. std::tie(found_inf_value, found_inf_bdim) = unwrapTensorAtLevel(found_inf.value(), cur_level);
  27891. }
  27892. batch_rule(self, grads, state_sums, state_steps, lr_value, lr_bdim, lr_decay, weight_decay, eps, maximize, grad_scale_value, grad_scale_bdim, found_inf_value, found_inf_bdim, out);
  27893. }
  27894. template <typename batch_rule_t, batch_rule_t batch_rule>
  27895. ::std::tuple<::std::vector<at::Tensor>,::std::vector<at::Tensor>,::std::vector<at::Tensor>> _fused_adagrad_tensor_lr_generated_plumbing(at::TensorList self, at::TensorList grads, at::TensorList state_sums, at::TensorList state_steps, const at::Tensor & lr, double lr_decay, double weight_decay, double eps, bool maximize, const ::std::optional<at::Tensor> & grad_scale, const ::std::optional<at::Tensor> & found_inf) {
  27896. c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);
  27897. auto maybe_layer = maybeCurrentDynamicLayer();
  27898. vmap_check_escaped(maybe_layer, "gen_vmap_plumbing");
  27899. int64_t cur_level = maybe_layer->layerId();
  27900. if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(grads, cur_level) && !isBatchedAtLevel(state_sums, cur_level) && !isBatchedAtLevel(state_steps, cur_level) && !isBatchedAtLevel(lr, cur_level) && !isBatchedAtLevel(grad_scale, cur_level) && !isBatchedAtLevel(found_inf, cur_level)) {
  27901. return at::_ops::_fused_adagrad_tensor_lr::call(self, grads, state_sums, state_steps, lr, lr_decay, weight_decay, eps, maximize, grad_scale, found_inf);
  27902. }
  27903. auto [lr_value, lr_bdim] = unwrapTensorAtLevel(lr, cur_level);
  27904. std::optional<Tensor> grad_scale_value;
  27905. std::optional<int64_t> grad_scale_bdim;
  27906. if (grad_scale) {
  27907. std::tie(grad_scale_value, grad_scale_bdim) = unwrapTensorAtLevel(grad_scale.value(), cur_level);
  27908. }
  27909. std::optional<Tensor> found_inf_value;
  27910. std::optional<int64_t> found_inf_bdim;
  27911. if (found_inf) {
  27912. std::tie(found_inf_value, found_inf_bdim) = unwrapTensorAtLevel(found_inf.value(), cur_level);
  27913. }
  27914. auto results = batch_rule(self, grads, state_sums, state_steps, lr_value, lr_bdim, lr_decay, weight_decay, eps, maximize, grad_scale_value, grad_scale_bdim, found_inf_value, found_inf_bdim);
  27915. return std::make_tuple(makeBatchedVector(std::get<0>(results), std::get<1>(results), cur_level), makeBatchedVector(std::get<2>(results), std::get<3>(results), cur_level), makeBatchedVector(std::get<4>(results), std::get<5>(results), cur_level));
  27916. }
  27917. }} // namespace at::functorch