package Ġw um o . sim . tensorflow . ops . gradients Ċ Ċ import Ġregister _ array _ grad Ċ import Ġregister _ control _ flow _ grad Ċ import Ġregister _ c udnn _ rnn _ grad Ċ import Ġregister _ data _ flow _ grad Ċ import Ġregister _ image _ grad Ċ import Ġregister _ manip _ grad Ċ import Ġregister _ math _ grad Ċ import Ġregister _ nn _ grad Ċ import Ġregister _ random _ grad Ċ import Ġregister _ state _ grad Ċ import Ġregister _ tensor _ array _ grad Ċ import Ġw um o . sim . tensorflow . core . Invalid DataType Exception Ċ import Ġw um o . sim . tensorflow . ops .* Ċ import Ġw um o . sim . tensorflow . ops . control _ flow _ ops . Control Flow Context Ċ import Ġw um o . sim . tensorflow . ops . control _ flow _ ops . Control Flow State Ċ import Ġw um o . sim . tensorflow . ops . control _ flow _ ops . control _ flow _ ops Ċ import Ġw um o . sim . tensorflow . ops . gradients . gradient _ ops . Aggregation Method . Add Aggregation Method Ċ import Ġw um o . sim . tensorflow . tf Ċ import Ġw um o . sim . tensorflow . types .* Ċ import Ġw um o . sim . tensorflow . util . attr Value Ċ import Ġw um o . sim . util . debug Ċ import Ġw um o . sim . util . empty Mutable Set Ċ import Ġw um o . sim . util . lazy Logger Ċ import Ġw um o . sim . util . t 2 Ċ import Ġjava . util .* Ċ Ċ object Ġgradient _ ops Ġ{ ĊĠ Ġval Ġlogger Ġby Ġlazy Logger () ĊĠĠĊĠ Ġinit Ġ{ ĊĠĠĠ Ġregister _ array _ grad () ĊĠĠĠ Ġregister _ control _ flow _ grad () ĊĠĠĠ Ġregister _ c udnn _ rnn _ grad () ĊĠĠĠ Ġregister _ data _ flow _ grad () ĊĠĠĠ Ġregister _ image _ grad () ĊĠĠĠ Ġregister _ manip _ grad () ĊĠĠĠ Ġregister _ math _ grad () ĊĠĠĠ Ġregister _ nn _ grad () ĊĠĠĠ Ġregister _ random _ grad () ĊĠĠĠ Ġregister _ state _ grad () ĊĠĠĠ Ġregister _ tensor _ array _ grad () ĊĠ Ġ} ĊĠĠĊĠ Ġinterface ĠAPI Ġ{ Ċ // ĠĠĠ Ġfun Ġgradients ( y : ĠOutput , Ġxs : ĠCollection < Output >): ĠList < Output > Ġ{ Ċ // ĠĠĠĠĠ Ġreturn Ġadd Symbolic Grad ients ( listOf ( y ), Ġxs . toList ()) Ċ // ĠĠĠ Ġ} Ċ // Ċ // ĠĠĠ Ġfun Ġgradients ( ys : ĠList < Output >, Ġxs : ĠCollection < Output >): ĠList < Output > Ġ{ Ċ // ĠĠĠĠĠ Ġreturn Ġadd Symbolic Grad ients ( ys , Ġxs . toList ()) Ċ // ĠĠĠ Ġ} ĊĠĠĠĠĊĠĠĠ Ġ/** ĊĠĠĠĠ Ġ* ĠConstructs Ġsymbolic Ġderivatives Ġof Ġsum Ġof Ġ` ys ` Ġw . r . t . Ġx Ġin Ġ` xs `. ĊĠĠĠĠ Ġ* ĊĠĠĠĠ Ġ* Ġ` ys ` Ġand Ġ` xs ` Ġare Ġeach Ġa Ġ` Tensor ` Ġor Ġa Ġlist Ġof Ġtensors . Ġ Ġ` grad _ ys ` ĊĠĠĠĠ Ġ* Ġis Ġa Ġlist Ġof Ġ` Tensor `, Ġholding Ġthe Ġgradients Ġreceived Ġby Ġthe ĊĠĠĠĠ Ġ* Ġ` ys `. ĠThe Ġlist Ġmust Ġbe Ġthe Ġsame Ġlength Ġas Ġ` ys `. ĊĠĠĠĠ Ġ* ĊĠĠĠĠ Ġ* Ġ` gradients ()` Ġadds Ġops Ġto Ġthe Ġgraph Ġto Ġoutput Ġthe Ġderivatives Ġof Ġ` ys ` Ġwith ĊĠĠĠĠ Ġ* Ġrespect Ġto Ġ` xs `. Ġ ĠIt Ġreturns Ġa Ġlist Ġof Ġ` Tensor ` Ġof Ġlength Ġ` len ( xs )` Ġwhere ĊĠĠĠĠ Ġ* Ġeach Ġtensor Ġis Ġthe Ġ` sum ( dy / dx )` Ġfor Ġy Ġin Ġ` ys `. ĊĠĠĠĠĠ ĊĠĠĠĠ Ġ* Ġ` grad _ ys ` Ġis Ġa Ġlist Ġof Ġtensors Ġof Ġthe Ġsame Ġlength Ġas Ġ` ys ` Ġthat Ġholds ĊĠĠĠĠ Ġ* Ġthe Ġinitial Ġgradients Ġfor Ġeach Ġy Ġin Ġ` ys `. Ġ ĠWhen Ġ` grad _ ys ` Ġis ĠNone , ĊĠĠĠ Ġwe Ġfill Ġin Ġa Ġtensor Ġof Ġ' 1 's Ġof Ġthe Ġshape Ġof Ġy Ġfor Ġeach Ġy Ġin Ġ` ys `. Ġ ĠA ĊĠĠĠ Ġuser Ġcan Ġprovide Ġtheir Ġown Ġinitial Ġ` grad _ ys ` Ġto Ġcompute Ġthe ĊĠĠĠ Ġderivatives Ġusing Ġa Ġdifferent Ġinitial Ġgradient Ġfor Ġeach Ġy Ġ( e . g ., Ġif ĊĠĠĠ Ġone Ġwanted Ġto Ġweight Ġthe Ġgradient Ġdifferently Ġfor Ġeach Ġvalue Ġin ĊĠĠĠ Ġeach Ġy ). ĊĠĠĠĠĊĠĠĠ Ġ` stop _ gradients ` Ġis Ġa Ġ` Tensor ` Ġor Ġa Ġlist Ġof Ġtensors Ġto Ġbe Ġconsidered Ġconstant ĊĠĠĠ Ġwith Ġrespect Ġto Ġall Ġ` xs `. ĠThese Ġtensors Ġwill Ġnot Ġbe Ġback propag ated Ġthrough , ĊĠĠĠ Ġas Ġthough Ġthey Ġhad Ġbeen Ġexplicitly Ġdisconnected Ġusing Ġ` stop _ gradient `. Ġ ĠA mong ĊĠĠĠ Ġother Ġthings , Ġthis Ġallows Ġcomputation Ġof Ġpartial Ġderivatives Ġas Ġopposed Ġto ĊĠĠĠ Ġtotal Ġderivatives . ĠFor Ġexample : ĊĠĠĠĠĊĠĠĠ Ġ``` python ĊĠĠĠ Ġa Ġ= Ġtf . constant ( 0 .) ĊĠĠĠ Ġb Ġ= Ġ 2 Ġ* Ġa ĊĠĠĠ Ġg Ġ= Ġtf . gradients ( a Ġ+ Ġb , Ġ[ a , Ġb ], Ġstop _ gradients =[ a , Ġb ]) ĊĠĠĠ Ġ``` ĊĠĠĠĠĊĠĠĠ ĠHere Ġthe Ġpartial Ġderivatives Ġ` g ` Ġevaluate Ġto Ġ`[ 1 . 0 , Ġ 1 . 0 ] `, Ġcompared Ġto Ġthe ĊĠĠĠ Ġtotal Ġderivatives Ġ` tf . gradients ( a Ġ+ Ġb , Ġ[ a , Ġb ]) `, Ġwhich Ġtake Ġinto Ġaccount Ġthe ĊĠĠĠ Ġinfluence Ġof Ġ` a ` Ġon Ġ` b ` Ġand Ġevaluate Ġto Ġ`[ 3 . 0 , Ġ 1 . 0 ] `. Ġ ĠNote Ġthat Ġthe Ġabove Ġis ĊĠĠĠ Ġequivalent Ġto : ĊĠĠĠĠĊĠĠĠ Ġ``` python ĊĠĠĠ Ġa Ġ= Ġtf . stop _ gradient ( tf . constant ( 0 . )) ĊĠĠĠ Ġb Ġ= Ġtf . stop _ gradient ( 2 Ġ* Ġa ) ĊĠĠĠ Ġg Ġ= Ġtf . gradients ( a Ġ+ Ġb , Ġ[ a , Ġb ]) ĊĠĠĠ Ġ``` ĊĠĠĠĠĊĠĠĠ Ġ` stop _ gradients ` Ġprovides Ġa Ġway Ġof Ġstopping Ġgradient Ġafter Ġthe Ġgraph Ġhas ĊĠĠĠ Ġalready Ġbeen Ġconstructed , Ġas Ġcompared Ġto Ġ` tf . stop _ gradient ` Ġwhich Ġis Ġused ĊĠĠĠ Ġduring Ġgraph Ġconstruction . Ġ ĠWhen Ġthe Ġtwo Ġapproaches Ġare Ġcombined , ĊĠĠĠ Ġback propagation Ġstops Ġat Ġboth Ġ` tf . stop _ gradient ` Ġnodes Ġand Ġnodes Ġin ĊĠĠĠ Ġ` stop _ gradients `, Ġwh iche ver Ġis Ġencountered Ġfirst . ĊĠĠĠĠĊĠĠĠ ĠAll Ġinteger Ġtensors Ġare Ġconsidered Ġconstant Ġwith Ġrespect Ġto Ġall Ġ` xs `, Ġas Ġif ĊĠĠĠ Ġthey Ġwere Ġincluded Ġin Ġ` stop _ gradients `. ĊĠĠĠĠĊĠĠĠ ĠArgs : ĊĠĠĠ Ġys : ĠA Ġ` Tensor ` Ġor Ġlist Ġof Ġtensors Ġto Ġbe Ġdifferent i ated . ĊĠĠĠ Ġxs : ĠA Ġ` Tensor ` Ġor Ġlist Ġof Ġtensors Ġto Ġbe Ġused Ġfor Ġdifferent iation . ĊĠĠĠ Ġgrad _ ys : ĠOptional . ĠA Ġ` Tensor ` Ġor Ġlist Ġof Ġtensors Ġthe Ġsame Ġsize Ġas ĊĠĠĠ Ġ` ys ` Ġand Ġholding Ġthe Ġgradients Ġcomputed Ġfor Ġeach Ġy Ġin Ġ` ys `. ĊĠĠĠ Ġname : ĠOptional Ġname Ġto Ġuse Ġfor Ġgrouping Ġall Ġthe Ġgradient Ġops Ġtogether . ĊĠĠĠ Ġdefaults Ġto Ġ' gradients '. ĊĠĠĠ Ġco locate _ gradients _ with _ ops : ĠIf ĠTrue , Ġtry Ġco loc ating Ġgradients Ġwith ĊĠĠĠ Ġthe Ġcorresponding Ġop . ĊĠĠĠ Ġgate _ gradients : ĠIf ĠTrue , Ġadd Ġa Ġtuple Ġaround Ġthe Ġgradients Ġreturned ĊĠĠĠ Ġfor Ġan Ġoperations . Ġ ĠThis Ġavoids Ġsome Ġrace Ġconditions . ĊĠĠĠ Ġaggregation _ method : ĠSpecifies Ġthe Ġmethod Ġused Ġto Ġcombine Ġgradient Ġterms . ĊĠĠĠ ĠAccept ed Ġvalues Ġare Ġconstants Ġdefined Ġin Ġthe Ġclass Ġ` Aggregation Method `. ĊĠĠĠ Ġstop _ gradients : ĠOptional . ĠA Ġ` Tensor ` Ġor Ġlist Ġof Ġtensors Ġnot Ġto Ġdifferent iate ĊĠĠĠ Ġthrough . ĊĠĠĠĠĊĠĠĠ ĠReturns : ĊĠĠĠ ĠA Ġlist Ġof Ġ` sum ( dy / dx )` Ġfor Ġeach Ġx Ġin Ġ` xs `. ĊĠĠĠĠ Ġ* ĊĠĠĠĠ Ġ* Ġ@ see Ġ" tensorflow . python . ops . gradients _ impl . gradients " ĊĠĠĠĠ Ġ*/ ĊĠĠĠ Ġfun Ġgradients ( ys : ĠList < Output >, ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġxs : ĠList < Output >, ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġgrad _ ys : ĠList < Output Like >? Ġ= Ġnull , ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġgate Grad ients : ĠBoolean Ġ= Ġfalse , ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġaggregation Method : ĠAg gregation Method Ġ= ĠAdd Aggregation Method , ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġco locate Gradient sWith Ops : ĠBoolean Ġ= Ġfalse , ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġname : ĠString Ġ= Ġ" gradients "): ĠList < Output Like ?> Ġ{ ĊĠĠĠĠĠ Ġval Ġto _ ops Ġ= Ġys . map To ( mutable Set Of ()) Ġ{ Ġit . op Ġ} ĊĠĠĠĠĠ Ġval Ġfrom _ ops Ġ= Ġxs . map To ( mutable Set Of ()) Ġ{ Ġit . op Ġ} ĊĠĠĠĠĠ Ġval Ġgrad _ ops Ġ= Ġgrad _ ys ?. map To ( mutable Set Of ()) Ġ{ Ġit . op Ġ} Ġ?: Ġempty Mutable Set < Op >() ĊĠĠĠĠĠĠĊĠĠĠĠĠ Ġval Ġgrad s Ġ= Ġtf . name Scope ( name , Ġto _ ops Ġ+ Ġfrom _ ops Ġ+ Ġgrad _ ops ) Ġ{ ĊĠĠĠĠĠĠĠ Ġ// ĠGet Ġa Ġuid Ġfor Ġthis Ġcall Ġto Ġgradients Ġthat Ġcan Ġbe Ġused Ġto Ġhelp ĊĠĠĠĠĠĠĠ Ġ// Ġcluster Ġops Ġfor Ġcompilation . ĊĠĠĠĠĠĠĠ Ġval Ġgradient UID Ġ= Ġtf . current Graph . unique Name (" uid ") ĊĠĠĠĠĠĠĠĠĊĠĠĠĠĠĠĠ Ġ// ĠThe Ġapproach Ġwe Ġtake Ġhere Ġis Ġas Ġfollows : ĠCreate Ġa Ġlist Ġof Ġall Ġops Ġin Ġthe ĊĠĠĠĠĠĠĠ Ġ// Ġsub graph Ġbetween Ġthe Ġys Ġand Ġxs . Ġ ĠVisit Ġthese Ġops Ġin Ġreverse Ġorder Ġof Ġids ĊĠĠĠĠĠĠĠ Ġ// Ġto Ġensure Ġthat Ġwhen Ġwe Ġvisit Ġan Ġop Ġthe Ġgradients Ġw . r . t Ġits Ġoutputs Ġhave ĊĠĠĠĠĠĠĠ Ġ// Ġbeen Ġcollected . Ġ ĠThen Ġaggregate Ġthese Ġgradients Ġif Ġneeded , Ġcall Ġthe Ġop 's ĊĠĠĠĠĠĠĠ Ġ// Ġgradient Ġfunction , Ġand Ġadd Ġthe Ġgenerated Ġgradients Ġto Ġthe Ġgradients Ġfor ĊĠĠĠĠĠĠĠ Ġ// Ġits Ġinput . ĊĠĠĠĠĠĠĠĠĊĠĠĠĠĠĠĠ Ġ// ĠInitialize Ġthe Ġpending Ġcount Ġfor Ġops Ġin Ġthe Ġconnected Ġsub graph Ġfrom Ġys ĊĠĠĠĠĠĠĠ Ġval Ġto Ops Ġ= Ġif Ġ( ys . size Ġ> Ġ 1 ) ĊĠĠĠĠĠĠĠĠĠ Ġys . map To ( mutable Set Of ()) Ġ{ Ġy Ġ-> ĊĠĠĠĠĠĠĠĠĠĠĠ Ġif Ġ( y . consum ers . isNotEmpty ()) Ġtf . identity ( y ). op Ġelse Ġy . op ĊĠĠĠĠĠĠĠĠĠ Ġ} ĊĠĠĠĠĠĠĠ Ġelse ĊĠĠĠĠĠĠĠĠĠ Ġto _ ops ĊĠĠĠĠĠĠĠĠĊĠĠĠĠĠĠĠ Ġ// Ġ` pending Counts ( op )` Ġis Ġa Ġcount - down Ġcounter Ġfor Ġthe Ġexpected Ġgradients Ġto Ġaccum ulate Ġfor Ġ` op `. ĠWhen ĊĠĠĠĠĠĠĠ Ġ// Ġ` pending Counts ( op )` Ġbecomes Ġzero , Ġwe Ġhave Ġcollected Ġall Ġthe Ġback propagation Ġgradients Ġfor Ġall Ġoutputs Ġof Ġ` op `. ĊĠĠĠĠĠĠĠ Ġval Ġ( pending Count , Ġloop State ) Ġ= Ġinitial Pending Counts ( to Ops , Ġfrom _ ops , Ġco locate Gradient sWith Ops ) ĊĠĠĠĠĠĠĠĠĊĠĠĠĠĠĠĠ Ġ// ĠIterate Ġover Ġthe Ġcollected Ġops . ĊĠĠĠĠĠĠĠ Ġ// ĊĠĠĠĠĠĠĠ Ġ// Ġgrad s : Ġop Ġ=> Ġlist Ġof Ġgradients Ġreceived Ġon Ġeach Ġoutput Ġendpoint Ġof Ġthe ĊĠĠĠĠĠĠĠ Ġ// Ġop . Ġ ĠThe Ġgradients Ġfor Ġeach Ġendpoint Ġare Ġinitially Ġcollected Ġas Ġa Ġlist . ĊĠĠĠĠĠĠĠ Ġ// ĠWhen Ġit Ġis Ġtime Ġto Ġcall Ġthe Ġop 's Ġgradient Ġfunction , Ġfor Ġeach Ġendpoint Ġwe ĊĠĠĠĠĠĠĠ Ġ// Ġaggregate Ġthe Ġlist Ġof Ġreceived Ġgradients Ġinto Ġa ĠAdd () ĠOperation Ġif Ġthere ĊĠĠĠĠĠĠĠ Ġ// Ġis Ġmore Ġthan Ġone . ĊĠĠĠĠĠĠĠ Ġval Ġgrad s Ġ= Ġmutable Map Of < Op , ĠMutable List < Mutable List < Output Like ? >> >() ĊĠĠĠĠĠĠĠĠĊĠĠĠĠĠĠĠ Ġ// ĠAdd Ġthe Ġinitial Ġgradients Ġfor Ġthe Ġys . ĊĠĠĠĠĠĠĠ Ġval Ġgrad _ ys Ġ= Ġinitial Grad ients ( ys , Ġgrad _ ys , Ġco locate Gradient sWith Ops , Ġgradient UID ) ĊĠĠĠĠĠĠĠĠĊĠĠĠĠĠĠĠ Ġys . as Sequence (). zip ( grad _ ys . as Sequence ()). forEach Ġ{ Ġ( y , Ġgrad _ y ) Ġ-> ĊĠĠĠĠĠĠĠĠĠ Ġset Gradient ( grad s , Ġy , Ġgrad _ y ) ĊĠĠĠĠĠĠĠ Ġ} ĊĠĠĠĠĠĠĠĠĊĠĠĠĠĠĠĠ Ġ// Ġ` ready Ops ` Ġkeeps Ġtrack Ġof Ġops Ġthat Ġhave Ġbeen Ġcompletely Ġprocessed . ĠWe Ġinitialize Ġit Ġwith Ġthe Ġdestination Ġops . ĊĠĠĠĠĠĠĠ Ġ// ĠWe Ġfilter Ġthe Ġdestination Ġops Ġbased Ġon Ġwhether Ġone Ġoutput Ġgradient Ġrelies Ġon Ġanother Ġoutput 's Ġgradient . ĊĠĠĠĠĠĠĠ Ġval Ġready Ops Ġ= ĠArray Deque < Op >( to Ops . filter Ġ{ Ġpending Count . get OrDefault ( it , Ġ 0 ) Ġ== Ġ 0 Ġ}) ĊĠĠĠĠĠĠĠĠĊĠĠĠĠĠĠĠ Ġloop State ?. let Ġ{ ĊĠĠĠĠĠĠĠĠĠ Ġit . process Unused Loop Ex its ( pending Count , Ġready Ops . toSet ()) ĊĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ. filter Ġ{ Ġis Train able ( it ) Ġ} ĊĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ. forEach Ġ{ Ġloop Exit Ġ-> ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġset Gradient ( grad s , Ġloop Exit , Ġloop State . zeros Like For Exit ( loop Exit )) ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġready Ops Ġ+= Ġloop Exit . op ĊĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ} ĊĠĠĠĠĠĠĠ Ġ} ĊĠĠĠĠĠĠĠĠĊĠĠĠĠĠĠĠ Ġ// ĠStop Ġops Ġform Ġthe Ġfront ier Ġof Ġthe Ġforward Ġgraph Ġbefore Ġwhich Ġback - propagation Ġshould Ġstop . ĠOps Ġin Ġthis Ġset Ġwill ĊĠĠĠĠĠĠĠ Ġ// Ġnot Ġbe Ġdifferent i ated . ĠThis Ġset Ġis Ġdefined Ġas Ġthe Ġsubset Ġof Ġ` source Ops ` Ġcontaining Ġops Ġthat Ġhave Ġno Ġpre decessor ĊĠĠĠĠĠĠĠ Ġ// Ġin Ġ` source Ops `. ĠAn Ġop Ġhas Ġpre decess ors Ġin Ġ` source Ops ` Ġif Ġand Ġonly Ġif Ġ` pending Counts ( op ) Ġ> Ġ 0 `. ĊĠĠĠĠĠĠĠ Ġval Ġstop Ops Ġ= Ġfrom _ ops . filter To ( mutable Set Of ()) Ġ{ ĊĠĠĠĠĠĠĠĠĠ Ġit . inputs . all Ġ{ Ġpending Count . get OrDefault ( it . op , Ġ 0 ) Ġ<= Ġ 0 Ġ} ĊĠĠĠĠĠĠĠ Ġ} ĊĠĠĠĠĠĠĠĠĊĠĠĠĠĠĠĠ Ġwhile Ġ( ready Ops . isNotEmpty ()) Ġ{ ĊĠĠĠĠĠĠĠĠĠ Ġ// Ġgenerate Ġgradient Ġsub graph Ġfor Ġop . ĊĠĠĠĠĠĠĠĠĠ Ġval Ġop Ġ= Ġready Ops . remove First () ĊĠĠĠĠĠĠĠĠĠ Ġmaybe Co locate With ( op , Ġco locate Gradient sWith Ops , Ġgradient UID ) Ġ{ ĊĠĠĠĠĠĠĠĠĠĠĠ Ġloop State ?. enter Gradient While Context ( op , Ġbefore Ġ= Ġtrue ) ĊĠĠĠĠĠĠĠĠĠĠĠ Ġval Ġop Grad ients Ġ= Ġaggregation Method . aggregate Grad ients ( grad s , Ġop , Ġgradient UID ) ĊĠĠĠĠĠĠĠĠĠĠĠ Ġloop State ?. exit Gradient While Context ( op , Ġbefore Ġ= Ġtrue ) ĊĠĠĠĠĠĠĠĠĠĠĠĠĊĠĠĠĠĠĠĠĠĠĠĠ Ġval Ġhas Output Grad ients Ġ= Ġop Grad ients . isNotEmpty () ĊĠĠĠĠĠĠĠĠĠĠĠ Ġval Ġgradient Function Ġ= Ġif Ġ( has Output Grad ients Ġ&& Ġop Ġ! in Ġstop Ops ) ĊĠĠĠĠĠĠĠĠĠĠĠĠĠ ĠRegistry [ op . op Type ] ĊĠĠĠĠĠĠĠĠĠĠĠ Ġelse Ġnull ĊĠĠĠĠĠĠĠĠĠĠĠ Ġloop State ?. enter Gradient While Context ( op , Ġbefore Ġ= Ġfalse ) ĊĠĠĠĠĠĠĠĠĠĠĠ Ġif Ġ( has Output Grad ients Ġ&& Ġgradient Function Ġ!= Ġnull ) Ġ{ ĊĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ// ĠNote Ġthat , Ġthe Ġgradient Ġaggregation Ġnot Ġcomputing Ġa Ġvalue Ġfor Ġthe Ġi 't h Ġoutput , Ġmeans Ġthat Ġthe Ġcost Ġdoes ĊĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ// Ġnot Ġdepend Ġon Ġoutput Ġi Ġand Ġtherefore Ġthe Ġgradient Ġwith Ġrespect Ġto Ġthat Ġoutput Ġis Ġ 0 . ĊĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġop Grad ients . with Index (). forEach Ġ{ Ġ( output Index , Ġgradient ) Ġ-> ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ// ĠOnly Ġfloating - point Ġoutputs Ġget Ġa Ġzero Ġgradient . ĠGradient Ġfunctions Ġshould Ġignore Ġthe Ġgradient Ġfor ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ// Ġother Ġoutputs . ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġval Ġoutput Ġ= Ġop . outputs [ output Index ] ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġif Ġ( gradient . isEmpty () Ġ&& Ġis Train able ( output )) ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġop Grad ients [ output Index ] Ġ= Ġmutable ListOf ( loop State ?. zeros Like ( op , Ġoutput Index ) ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ?: ĠControl Flow Context . zeros Like Outside Loop ( op , Ġoutput Index )) ĊĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ} ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĊĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ// ĠCompute Ġthe Ġactual Ġop Ġgradients . ĊĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġtf . name Scope ("${ op . name }_ grad ") Ġ{ ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ// ĠTODO : Ġ[ CONTEXT ] ĠAdd Ġsupport Ġfor Ġoriginal Ġop Ġcontext . ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġval Ġoutput Grad ients Ġ= Ġop Grad ients . map Ġ{ Ġit . first OrNull () Ġ} ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġvar Ġinput Grad ients Ġ= Ġmaybe Compile ( name , Ġop ) Ġ{ Ġgradient Function ( op , Ġoutput Grad ients ) Ġ} ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġif Ġ( gate Grad ients Ġ&& Ġinput Grad ients . count Ġ{ Ġit Ġ!= Ġnull Ġ} Ġ> Ġ 1 ) ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġtf . device ( dev Ġ= Ġ"") Ġ{ ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġtf . co locate With For Gradient ( mutable Set Of (), Ġgradient UID , Ġignore Existing Ġ= Ġtrue ) Ġ{ ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġinput Grad ients Ġ= Ġtf . tuple ( input Grad ients ) ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ} ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ} ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġval Ġn In p Ġ= Ġop . inputs . size ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġval Ġn Gr d Ġ= Ġinput Grad ients . size ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġassert ( n In p Ġ== Ġn Gr d ) Ġ{ Ġ" Grad ients Ġsize Ġ($ n Gr d ) Ġfor Ġop Ġ'$ op ' Ġdoes Ġnot Ġmatch Ġinputs Ġsize Ġ($ n In p )." Ġ} ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġlog Grad ients ( op , Ġoutput Grad ients , Ġinput Grad ients ) ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ// ĠTODO : ĠReport Ġsomehow Ġthe Ġnon - different iable Ġops Ġin Ġthe Ġgraph . ĠThis Ġis Ġcurrently Ġhard Ġto Ġdebug . ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġop . inputs . as Sequence (). zip ( input Grad ients . as Sequence ()) ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ. forEach Ġ{ Ġ( t In , Ġin Grad ) Ġ-> ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġif Ġ( in Grad Ġ== Ġnull ) Ġreturn @ forEach ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġif Ġ( in Grad Ġis ĠOutput Ġ&& Ġt In . dataType Ġ!= ĠRESOURCE ) ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġin Grad . set Shape ( t In . shape ) ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġset Gradient ( grad s , Ġt In , Ġin Grad ) ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ} ĊĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ} ĊĠĠĠĠĠĠĠĠĠĠĠ Ġ} ĊĠĠĠĠĠĠĠĠĠĠĠ Ġloop State ?. exit Gradient While Context ( op , Ġbefore Ġ= Ġfalse ) ĊĠĠĠĠĠĠĠĠĠ Ġ} ĊĠĠĠĠĠĠĠĠĠĠĊĠĠĠĠĠĠĠĠĠ Ġ// ĠUpdate Ġthe Ġpending Ġcounts Ġfor Ġthe Ġinputs Ġof Ġ` op ` Ġand Ġenqueue Ġready Ġops . ĊĠĠĠĠĠĠĠĠĠ Ġop . inputs . forEach Ġ{ Ġinput Ġ-> ĊĠĠĠĠĠĠĠĠĠĠĠ Ġval Ġcount Ġ= Ġpending Count . compute ( input . op ) Ġ{ Ġ_, Ġcount Ġ-> ĊĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ( count Ġ?: Ġ 0 ) Ġ- Ġ 1 ĊĠĠĠĠĠĠĠĠĠĠĠ Ġ} !! ĊĠĠĠĠĠĠĠĠĠĠĠ Ġvar Ġready Ġ= Ġcount Ġ== Ġ 0 ĊĠĠĠĠĠĠĠĠĠĠĠ Ġif Ġ( loop State Ġ!= Ġnull Ġ&& Ġ! ready ) ĊĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġready Ġ= Ġcount Ġ> Ġ 0 Ġ&& Ġcontrol _ flow _ ops . is Loop Switch ( input . op ) ĊĠĠĠĠĠĠĠĠĠĠĠ Ġif Ġ( ready ) Ġ{ ĊĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġif Ġ( control _ flow _ ops . is Loop Exit ( input . op )) Ġ{ ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ// ĠIf Ġ` input ` Ġis Ġan Ġexit Ġwithout Ġreal Ġgradient , Ġdefer Ġprocessing Ġthem . ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġloop State ?. get Gradient Loop State ( input . op , Ġbefore Ġ= Ġfalse )?. let Ġ{ Ġgrad State Ġ-> ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġgrad State . deferred Ex its Ġ+= Ġinput ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġgrad State . pending Ex its Count Ġ-= Ġ 1 ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġif Ġ( grad State . pending Ex its Count Ġ== Ġ 0 ) Ġ{ ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ// ĠWe Ġnow Ġhave Ġall Ġthe Ġexits Ġand Ġso Ġwe Ġprocess Ġthem . ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġvar Ġhas Real Gradient Ġ= Ġfalse ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġgrad State . deferred Ex its . forEach Ġ{ Ġexit Ġ-> ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġif Ġ( grad s [ exit . op ]?. any Ġ{ Ġit . any Ġ{ Ġit Ġ!= Ġnull Ġ} Ġ} Ġ== Ġtrue ) Ġ{ ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġhas Real Gradient Ġ= Ġtrue ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġready Ops Ġ+= Ġexit . op ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ} Ġelse ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġgrad State . unused Ex its Ġ+= Ġexit ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ} ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġif Ġ( has Real Gradient ) ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ// ĠFor Ġan Ġunused Ġexit , Ġif Ġit Ġhas Ġfloating - point Ġoutputs , Ġwe Ġback - propagate Ġa Ġzero Ġgradient . ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ// ĠOtherwise , Ġwe Ġjust Ġignore Ġit . ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġgrad State . unused Ex its . forEach Ġ{ Ġexit Ġ-> ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġif Ġ( is Train able ( exit )) ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġset Gradient ( grad s , Ġexit , Ġloop State . zeros Like For Exit ( exit )) ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġready Ops Ġ+= Ġexit . op ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ} ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġelse ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ// ĠAll Ġexits Ġare Ġ" unused " Ġand Ġso Ġwe Ġuse Ġ` null ` Ġas Ġthe Ġgradient . ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġgrad State . unused Ex its . forEach Ġ{ Ġready Ops Ġ+= Ġit . op Ġ} ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ} ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ} ĊĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ} Ġelse ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġready Ops Ġ+= Ġinput . op ĊĠĠĠĠĠĠĠĠĠĠĠ Ġ} ĊĠĠĠĠĠĠĠĠĠ Ġ} ĊĠĠĠĠĠĠĠ Ġ} ĊĠĠĠĠĠĠĠĠĊĠĠĠĠĠĠĠ Ġloop State ?. post Process () ĊĠĠĠĠĠĠĠ Ġgrad s ĊĠĠĠĠĠ Ġ} ĊĠĠĠĠĠĠĊĠĠĠĠĠ Ġ// ĠCollect Ġthe Ġaggregated Ġgradients Ġfor Ġthe Ġrequested Ġtensors Ġand Ġreturn Ġthem . ĊĠĠĠĠĠ Ġreturn Ġxs . map Ġ{ Ġx Ġ-> ĊĠĠĠĠĠĠĠ Ġval Ġgradients Ġ= Ġgrad s [ x . op ]?. get ( x . value Index ) ĊĠĠĠĠĠĠĠ Ġif Ġ( gradients Ġ!= Ġnull Ġ&& Ġgradients . size Ġ> Ġ 1 ) ĊĠĠĠĠĠĠĠĠĠ Ġthrow ĠIllegalArgumentException (" The Ġgradients Ġshould Ġhave Ġbeen Ġaggregated Ġby Ġnow .") ĊĠĠĠĠĠĠĠ Ġgradients ?. first OrNull () ĊĠĠĠĠĠ Ġ} ĊĠĠĠ Ġ} ĊĠĠĠĠĊĠ Ġ} ĊĠĠĊĠ Ġ/** ĠIf Ġ` co locate Gradient sWith Ops ` Ġis Ġ` true `, Ġthen Ġall Ġops Ġcreated Ġwithin Ġ` block ` Ġwill Ġbe Ġco located Ġwith Ġ` op `. ĊĠĠ Ġ* ĊĠĠ Ġ* Ġ@ param Ġ Ġop ĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ ĠOp Ġto Ġmaybe Ġco locate Ġwith . ĊĠĠ Ġ* Ġ@ param Ġ Ġco locate Gradient sWith Ops ĠBoolean Ġvalue Ġindicating Ġwhether Ġto Ġco locate Ġthe Ġgradient Ġops Ġwith Ġthe Ġoriginal ĊĠĠ Ġ* ĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġops . ĊĠĠ Ġ* Ġ@ param Ġ Ġgradient UID ĠĠĠĠĠĠĠĠĠĠĠĠĠ ĠUnique Ġidentifier Ġwithin Ġthe Ġgraph Ġindicating Ġwhich Ġinvocation Ġof Ġgradients Ġis ĊĠĠ Ġ* ĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġbeing Ġexecuted . ĠUsed Ġto Ġcluster Ġops Ġfor Ġcompilation . ĊĠĠ Ġ* Ġ@ param Ġ Ġblock ĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ ĠBlock Ġof Ġcode Ġto Ġexecute Ġusing Ġthe Ġspecified Ġco location Ġops . ĊĠĠ Ġ* Ġ@ return ĠReturn Ġvalue Ġof Ġ` block `. ĊĠĠ Ġ*/ ĊĠ Ġprivate Ġfun Ġ< R > Ġmaybe Co locate With ( op : ĠOp , ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġco locate Gradient sWith Ops : ĠBoolean , ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġgradient UID : ĠString , ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġblock : Ġ() Ġ-> ĠR ): ĠR Ġ= ĊĠĠĠĠĠ Ġif Ġ( co locate Gradient sWith Ops ) ĊĠĠĠĠĠĠĠ Ġtf . co locate With For Gradient ( mutable Set Of ( op ), Ġgradient UID , Ġblock Ġ= Ġblock ) ĊĠĠĠĠĠ Ġelse ĊĠĠĠĠĠĠĠ Ġblock () ĊĠĠĊĠ Ġ/** ĠIf Ġthe Ġop Ġwas Ġmarked Ġas Ġcompiled , Ġthis Ġfunction Ġcompiles Ġthe Ġcalculation Ġin Ġ` gradient Function ` Ġ( using ĠX LA ) Ġand ĊĠĠ Ġ* Ġreturns Ġthe Ġresult Ġof Ġ` gradient Function `. ĠOtherwise , Ġit Ġsimply Ġreturns Ġthe Ġresult Ġof Ġ` gradient Function `. ĊĠĠ Ġ* ĊĠĠ Ġ* Ġ@ param Ġ Ġname Scope ĠĠĠĠĠĠĠ ĠName Ġnative Ġto Ġuse Ġfor Ġthe Ġgradient Ġops . ĊĠĠ Ġ* Ġ@ param Ġ Ġop ĠĠĠĠĠĠĠĠĠĠĠĠĠĠ ĠOp Ġwhose Ġgradients Ġare Ġbeing Ġcomputed . ĊĠĠ Ġ* Ġ@ param Ġ Ġgradient Function ĠFunction Ġthat Ġcomputes Ġthe Ġgradients Ġfor Ġ` op `. ĊĠĠ Ġ* Ġ@ return ĠCreated Ġgradients Ġop . ĊĠĠ Ġ*/ ĊĠ Ġprivate Ġfun Ġmaybe Compile ( name Scope : ĠString , Ġop : ĠOp , Ġgrad Fn : Ġ() Ġ-> ĠList < Output Like ?> ): ĠList < Output Like ?> Ġ{ ĊĠĠĠ Ġ// ĠTODO : Ġ[ FUNCTION AL ] ĠAdd Ġextra Ġ' func ' Ġargument . ĊĠĠĠ Ġval Ġclean Name Scope Ġ= Ġname Scope . remove Suffix ("/ "). replace ('/', Ġ'_ ') ĊĠĠĠ Ġreturn Ġtry Ġ{ ĊĠĠĠĠĠ Ġval Ġx la Compile Ġ= Ġop . attr Bool ("_ X la Compile ") ĊĠĠĠĠĠ Ġif Ġ(! x la Compile ) ĊĠĠĠĠĠĠĠ Ġgrad Fn () ĊĠĠĠĠĠ Ġelse Ġ{ ĊĠĠĠĠĠĠĠ Ġval Ġx la Separ ate Compile Gradient Ġ= Ġop . attr Bool ("_ X la Separ ate Compiled Grad ients ") ĊĠĠĠĠĠĠĠ Ġval Ġx la Scope Ġ= Ġop . attr String ("_ X la Scope ") ĊĠĠĠĠĠĠĠ Ġ// ĠIf Ġthe Ġgradients Ġare Ġsupposed Ġto Ġbe Ġcompiled Ġseparately , Ġwe Ġgive Ġthem Ġan Ġ'_ X la Scope ' Ġname Ġthat Ġis Ġbased Ġon ĊĠĠĠĠĠĠĠ Ġ// Ġthe Ġname _ scope Ġof Ġthe Ġgradients . ĠOtherwise , Ġthey Ġjust Ġinherit Ġthe Ġexisting Ġ'_ X la Scope ' Ġname , Ġwhich Ġlets Ġthem ĊĠĠĠĠĠĠĠ Ġ// Ġbe Ġmerged Ġtogether Ġwith Ġthe Ġnon - gradient Ġcomputation . ĊĠĠĠĠĠĠĠ Ġval Ġx la Grad ients Scope Ġ= Ġif Ġ( x la Separ ate Compile Gradient ) Ġ"${ x la Scope }_ grad _$ clean Name Scope " Ġelse Ġx la Scope ĊĠĠĠĠĠĠĠ Ġtf . attr Scope ("_ X la Compile " Ġto Ġattr Value ( x la Compile ), Ġ"_ X la Scope " Ġto Ġattr Value ( x la Grad ients Scope )) Ġ{ ĊĠĠĠĠĠĠĠĠĠ Ġgrad Fn () ĊĠĠĠĠĠĠĠ Ġ} ĊĠĠĠĠĠ Ġ} ĊĠĠĠ Ġ} Ġcatch Ġ( e : ĠException ) Ġ{ ĊĠĠĠĠĠ Ġgrad Fn () ĊĠĠĠ Ġ} ĊĠ Ġ} ĊĠĠĊĠ Ġ/** ĠComputes Ġinitial Ġvalues Ġfor Ġthe Ġprovided Ġgradients , Ġand Ġchecks Ġwhether Ġtheir Ġdata Ġtypes Ġare Ġcorrect . ĊĠĠ Ġ* ĊĠĠ Ġ* Ġ@ param Ġ Ġys ĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ ĠSequence Ġcontaining Ġthe Ġvariables Ġcorresponding Ġto Ġ` d ys `. ĊĠĠ Ġ* Ġ@ param Ġ Ġd ys ĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ ĠSequence Ġcontaining Ġtensor Ġgradients . ĊĠĠ Ġ* Ġ@ param Ġ Ġco locate Gradient sWith Ops ĠBoolean Ġvalue Ġindicating Ġwhether Ġto Ġco locate Ġthe Ġgradient Ġops Ġwith Ġthe Ġoriginal ĊĠĠ Ġ* ĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġops . ĊĠĠ Ġ* Ġ@ param Ġ Ġgradient UID ĠĠĠĠĠĠĠĠĠĠĠĠĠ ĠUnique Ġidentifier Ġwithin Ġthe Ġgraph Ġindicating Ġwhich Ġinvocation Ġof Ġgradients Ġis ĊĠĠ Ġ* ĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġbeing Ġexecuted . ĠUsed Ġto Ġcluster Ġops Ġfor Ġcompilation . ĊĠĠ Ġ* Ġ@ return ĠSequence Ġcontaining Ġthe Ġdefault Ġgradient Ġvalues . ĊĠĠ Ġ* Ġ@ throws ĠInvalid DataType Exception ĠIf Ġthe Ġgradient Ġtensor Ġdata Ġtypes Ġare Ġnot Ġcompatible Ġwith Ġthe Ġinput Ġdata Ġtypes . ĊĠĠ Ġ* ĊĠĠ Ġ* Ġ@ see Ġ" tensorflow . python . ops . gradients _ impl ._ Default Grad Ys " ĊĠĠ Ġ*/ ĊĠ Ġprivate Ġfun Ġinitial Grad ients ( ys : ĠList < Output >, ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġgrad _ ys : ĠList < Output Like ?> ?, ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġco locate Gradient sWith Ops : ĠBoolean , ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġgradient UID : ĠString ): ĠList < Output Like > Ġ{ ĊĠĠĠ Ġval Ġgrad _ ys Ġ= Ġgrad _ ys Ġ?: ĠList ( ys . size ) Ġ{ Ġnull Ġ} ĊĠĠĠ Ġif Ġ( grad _ ys . size Ġ!= Ġys . size ) ĊĠĠĠĠĠ Ġthrow ĠIllegalArgumentException (" Passed Ġ${ grad _ ys . size } Ġgrad _ ys Ġfor Ġ${ ys . size } Ġys ") ĊĠĠĠ Ġreturn Ġys . as Sequence (). zip ( grad _ ys . as Sequence ()). with Index (). map Ġ{ Ġ( index , Ġvalue ) Ġ-> ĊĠĠĠĠĠ Ġval Ġ( y , Ġgrad _ y ) Ġ= Ġvalue ĊĠĠĠĠĠ Ġif Ġ( grad _ y Ġ== Ġnull ) Ġ{ ĊĠĠĠĠĠĠĠ Ġif Ġ( y . dataType . is Complex ) ĊĠĠĠĠĠĠĠĠĠ Ġthrow ĠInvalid DataType Exception ( ĊĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ" Grad ients Ġof Ġcomplex Ġtensors Ġmust Ġset Ġ' gradients ' Ġ( variable . dataType Ġ= Ġ'${ y . dataType } '). ") ĊĠĠĠĠĠĠĠ Ġtf . fill ( tf . shape ( y ), ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġtf . const ( y . dataType , Ġ 1 , Ġname Ġ= Ġ" grad _ ys _$ index ")) ĊĠĠĠĠĠ Ġ} Ġelse Ġ{ ĊĠĠĠĠĠĠĠ Ġwhen Ġ{ ĊĠĠĠĠĠĠĠĠĠ Ġy . dataType . is Floating Point Ġ|| Ġy . dataType . is Integer Ġ-> ĊĠĠĠĠĠĠĠĠĠĠĠ Ġif Ġ(! grad _ y . dataType . is Floating Point Ġ&& Ġ! grad _ y . dataType . is Integer ) ĊĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġthrow ĠInvalid DataType Exception ( ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ" Gradient Ġdata Ġtype Ġ'${ grad _ y . dataType }' Ġgenerated Ġfor Ġreal Ġor Ġinteger - valu ed Ġtensor Ġ'$ y ' Ġwith Ġdata Ġtype Ġ" Ġ+ ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ"' ${ y . dataType }' Ġmust Ġbe Ġreal Ġor Ġinteger .") ĊĠĠĠĠĠĠĠĠĠ Ġy . dataType . is Complex Ġ-> ĊĠĠĠĠĠĠĠĠĠĠĠ Ġif Ġ(! grad _ y . dataType . is Complex ) ĊĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġthrow ĠInvalid DataType Exception ( ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ" Gradient Ġdata Ġtype Ġ'${ grad _ y . dataType }' Ġgenerated Ġfor Ġcomplex - valu ed Ġtensor Ġ'$ y ' Ġwith Ġdata Ġtype Ġ" Ġ+ ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ"' ${ y . dataType }' Ġmust Ġbe Ġcomplex .") ĊĠĠĠĠĠĠĠĠĠ Ġelse Ġ-> Ġ{ ĊĠĠĠĠĠĠĠĠĠĠĠ Ġthrow ĠInvalid DataType Exception ( ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ" Tensor Ġ'$ y ' Ġwith Ġdata Ġtype Ġ'${ y . dataType }' Ġmust Ġbe Ġnumeric Ġin Ġorder Ġto Ġobtain Ġa Ġdefault Ġgradient .") ĊĠĠĠĠĠĠĠĠĠ Ġ} ĊĠĠĠĠĠĠĠ Ġ} ĊĠĠĠĠĠĠĠ Ġ// ĠCreate Ġa Ġgradients Ġtensor Ġin Ġthe Ġname Ġnative Ġof Ġthe Ġgradients . ĠThis Ġis Ġrequired Ġin Ġorder Ġfor Ġtensor Ġarrays ĊĠĠĠĠĠĠĠ Ġ// Ġto Ġidentify Ġwhich Ġgradient Ġcall Ġa Ġgradient Ġvalue Ġis Ġcoming Ġfrom . ĊĠĠĠĠĠĠĠ Ġwhen Ġ( grad _ y ) Ġ{ ĊĠĠĠĠĠĠĠĠĠ Ġis ĠOutput Ġ-> Ġtf . identity ( grad _ y , Ġname Ġ= Ġ" grad _ ys _$ index ") ĊĠĠĠĠĠĠĠĠĠ Ġis ĠIndexed S lices Ġ-> ĊĠĠĠĠĠĠĠĠĠĠĠ ĠIndexed S lices ( tf . identity ( grad _ y . indices , Ġname Ġ= Ġ" grad _ ys _${ index }_ indices "), ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġtf . identity ( grad _ y . values , Ġname Ġ= Ġ" grad _ ys _${ index }_ values "), ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġgrad _ y . dense Shape . let Ġ{ ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġif Ġ( it Ġ== Ġnull ) Ġit ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġelse Ġtf . identity ( it , Ġname Ġ= Ġ" grad _ ys ${ index }_ shape ") ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ}) ĊĠĠĠĠĠĠĠĠĠ Ġis ĠSparse Output Ġ-> ĊĠĠĠĠĠĠĠĠĠĠĠ ĠSparse Output ( tf . identity ( grad _ y . indices , Ġname Ġ= Ġ" grad _ ys _${ index }_ indices "), ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġtf . identity ( grad _ y . values , Ġname Ġ= Ġ" grad _ ys _${ index }_ values "), ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġgrad _ y . dense Shape . let Ġ{ ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġif Ġ( it Ġ== Ġnull ) Ġit ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġelse Ġtf . identity ( it , Ġname Ġ= Ġ" grad _ ys ${ index }_ shape ") ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ}) ĊĠĠĠĠĠĠĠ Ġ} ĊĠĠĠĠĠ Ġ} ĊĠĠĠ Ġ}. toList () ĊĠ Ġ} ĊĠĠĊĠ Ġ/** ĠAdds Ġthe Ġprovided Ġ` gradient ` Ġto Ġthe Ġsequence Ġof Ġ` output `' s Ġgradients Ġthat Ġhave Ġbeen Ġcollected Ġso Ġfar . ĊĠĠ Ġ* ĊĠĠ Ġ* Ġ@ param Ġ Ġgradients ĠMap Ġwhere Ġthe Ġcollected Ġgradients Ġare Ġstored . ĊĠĠ Ġ* Ġ@ param Ġ Ġoutput ĠĠĠ ĠOp Ġoutput Ġwhose Ġgradient Ġis Ġprovided . ĊĠĠ Ġ* Ġ@ param Ġ Ġgradient Ġ ĠGradient Ġof Ġ` output ` Ġto Ġadd Ġto Ġthe Ġcollected Ġgradients . ĊĠĠ Ġ*/ ĊĠ Ġprivate Ġfun Ġset Gradient ( ĊĠĠĠĠĠ Ġgradients : ĠMutable Map < Op , ĠMutable List < Mutable List < Output Like ? >> >, ĊĠĠĠĠĠ Ġoutput : ĠOutput , ĊĠĠĠĠĠ Ġgradient : ĠOutput Like ĊĠ Ġ) Ġ{ ĊĠĠĠ Ġval Ġop Grad ients Ġ= Ġgradients . get Or Put ( output . op ) Ġ{ Ġoutput . op . outputs . map To ( mutable ListOf ()) Ġ{ Ġmutable ListOf < Output Like ?> () Ġ} Ġ} ĊĠĠĠ Ġif Ġ( control _ flow _ ops . is Loop Switch ( output . op )) ĊĠĠĠĠĠ Ġop Grad ients [ output . value Index ] Ġ= Ġmutable ListOf ( gradient ) ĊĠĠĠ Ġelse ĊĠĠĠĠĠ Ġop Grad ients [ output . value Index ]. add ( gradient ) ĊĠ Ġ} ĊĠĠĊĠ Ġ/** ĠLogs Ġthe Ġinput Ġand Ġoutput Ġgradients Ġof Ġthe Ġprovided Ġop . ĊĠĠ Ġ* ĊĠĠ Ġ* Ġ@ param Ġ Ġop ĠĠĠĠĠĠĠĠĠĠĠĠĠ ĠOp . ĊĠĠ Ġ* Ġ@ param Ġ Ġoutput Grad ients ĠOutput Ġgradients Ġof Ġop . ĊĠĠ Ġ* Ġ@ param Ġ Ġinput Grad ients Ġ ĠInput Ġgradients Ġof Ġop . ĊĠĠ Ġ*/ ĊĠ Ġprivate Ġfun Ġlog Grad ients ( op : ĠOp , Ġoutput Grad ients : ĠList < Output Like ? >, Ġinput Grad ients : ĠList < Output Like ?> ) Ġ{ ĊĠĠĠ Ġlogger . debug Ġ{ Ġ" Grad ients Ġfor Ġop Ġ'${ op . name } ': " Ġ} ĊĠĠĠ Ġlogger . debug Ġ{ Ġ" Ġ Ġin Ġ Ġ--> Ġ${ output Grad ients . filter Ġ{ Ġit Ġ!= Ġnull Ġ}. join ToString (", Ġ") Ġ{ Ġit !!. name Ġ}}" Ġ} ĊĠĠĠ Ġlogger . debug Ġ{ Ġ" Ġ Ġout Ġ--> Ġ${ input Grad ients . filter Ġ{ Ġit Ġ!= Ġnull Ġ}. join ToString (", Ġ") Ġ{ Ġit !!. name Ġ}}" Ġ} ĊĠ Ġ} ĊĠĠĊĠ Ġ/** ĠInitializes Ġthe Ġback - propagation Ġinput Ġcounts Ġfor Ġops Ġbetween Ġtwo Ġsets Ġof Ġops . ĊĠĠ Ġ* ĊĠĠ Ġ* Ġ' output Map ( op )' Ġindicates Ġthe Ġnumber Ġof Ġback - propagation Ġinputs Ġto Ġthis Ġop . ĊĠĠ Ġ* ĊĠĠ Ġ* Ġ@ param Ġ Ġto Ops ĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ ĠSet Ġof Ġsource Ġops . ĊĠĠ Ġ* Ġ@ param Ġ Ġfrom Ops ĠĠĠĠĠĠĠĠĠĠ ĠSet Ġof Ġdestination Ġops . ĊĠĠ Ġ* Ġ@ param Ġ Ġco locate Gradient sWith Ops ĠBoolean Ġvalue Ġindicating Ġwhether Ġto Ġco locate Ġthe Ġgradient Ġops Ġwith Ġthe Ġoriginal ĊĠĠ Ġ* ĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġops . ĊĠĠ Ġ* Ġ@ return ĠTuple Ġcontaining : Ġ( 1 ) ĠMap Ġfrom Ġop Ġto Ġthe Ġnumber Ġof Ġback - propagation Ġinputs Ġto Ġthis Ġop , Ġand Ġ( 2 ) Ġa Ġcontrol ĊĠĠ Ġ* ĠĠĠĠĠĠĠĠ Ġflow Ġgradient Ġstate Ġobject Ġwhich Ġis Ġnot Ġ` None ` Ġif Ġthe Ġops Ġbetween Ġ` sources ` Ġand Ġ` dest inations ` Ġcontain ĊĠĠ Ġ* ĠĠĠĠĠĠĠĠ Ġcontrol Ġflow Ġloops . ĊĠĠ Ġ*/ ĊĠ Ġprivate Ġfun Ġinitial Pending Counts ( ĊĠĠĠĠĠ Ġto Ops : ĠSet < Op >, ĊĠĠĠĠĠ Ġfrom Ops : ĠSet < Op >, ĊĠĠĠĠĠ Ġco locate Gradient sWith Ops : ĠBoolean ): Ġt 2 < Mutable Map < Op , ĠInt >, ĠControl Flow State ?> Ġ{ ĊĠĠĠ Ġ// Mark Ġreachable Ġops Ġfrom Ġfrom _ ops . ĊĠĠĠ Ġval Ġreached Ops Ġ= Ġmutable Set Of < Op >() ĊĠĠĠ Ġmark Reached Ops ( from Ops , Ġreached Ops ) ĊĠĠĠ Ġ// ĠX Ġin Ġreached _ ops Ġiff ĠX Ġis Ġreachable Ġfrom Ġfrom _ ops Ġby Ġa Ġpath Ġof Ġzero Ġor Ġmore ĊĠĠĠ Ġ// Ġback propag atable Ġtensors . ĊĠĠĠĠĊĠĠĠ Ġ// ĠMark Ġops Ġbetween Ġ` sources ' Ġand Ġ' dest inations ' ĊĠĠĠ Ġval Ġbetween Ops Ġ= Ġmutable Set Of < Op >() ĊĠĠĠ Ġval Ġbetween Op List Ġ= Ġmutable ListOf < Op >() ĊĠĠĠ Ġval Ġbetween Queue Ġ= ĠArray Deque < Op >( to Ops ) ĊĠĠĠ Ġwhile Ġ( between Queue . isNotEmpty ()) Ġ{ ĊĠĠĠĠĠ Ġval Ġop Ġ= Ġbetween Queue . remove First () ĊĠĠĠĠĠ Ġif Ġ( op Ġin Ġreached Ops ) Ġ{ ĊĠĠĠĠĠĠĠ Ġbetween Ops Ġ+= Ġop ĊĠĠĠĠĠĠĠ Ġbetween Op List Ġ+= Ġop ĊĠĠĠĠĠĠĠ Ġreached Ops Ġ-= Ġop Ġ// Clear Ġthe Ġboolean Ġso Ġwe Ġwon 't Ġadd Ġthe Ġinputs Ġagain . ĊĠĠĠĠĠĠĠ Ġop . inputs . forEach Ġ{ Ġbetween Queue . add Last ( it . op ) Ġ} ĊĠĠĠĠĠ Ġ} ĊĠĠĠ Ġ} ĊĠĠĠ Ġ// ĠX Ġin Ġbetween _ ops Ġiff ĠX Ġis Ġon Ġa Ġpath Ġof Ġzero Ġor Ġmore Ġback propag atable Ġtensors ĊĠĠĠ Ġ// Ġbetween Ġfrom _ ops Ġand Ġto _ ops ĊĠĠĠ Ġval Ġloop State Ġ= ĠControl Flow State . maybe Create ( between Ops , Ġbetween Op List , Ġco locate Gradient sWith Ops ) ĊĠĠĠĠĊĠĠĠ Ġ// ĠInitialize Ġpending Ġcount Ġfor Ġbetween Ġops . ĊĠĠĠ Ġval Ġpending Count Ġ= Ġmutable Map Of < Op , ĠInt >() ĊĠĠĠ Ġbetween Op List . as Sequence () ĊĠĠĠĠĠĠĠ Ġ. flatMap Ġ{ Ġit . inputs . as Sequence () Ġ} ĊĠĠĠĠĠĠĠ Ġ. map Ġ{ Ġit . op Ġ} ĊĠĠĠĠĠĠĠ Ġ. filter Ġ{ Ġit Ġin Ġbetween Ops Ġ} ĊĠĠĠĠĠĠĠ Ġ. forEach Ġ{ ĊĠĠĠĠĠĠĠĠĠ Ġpending Count . compute ( it ) Ġ{ Ġ_, Ġcount Ġ-> ĊĠĠĠĠĠĠĠĠĠĠĠ Ġ( count Ġ?: Ġ 0 ) Ġ+ Ġ 1 ĊĠĠĠĠĠĠĠĠĠ Ġ} ĊĠĠĠĠĠĠĠ Ġ} ĊĠĠĠ Ġreturn Ġt 2 ( pending Count , Ġloop State ) ĊĠ Ġ} ĊĠĠĊĠ Ġ/** ĊĠĠ Ġ* ĠMark Ġall Ġops Ġreached Ġfrom Ġ" from _ ops ". ĊĠĠ Ġ*/ ĊĠ Ġprivate Ġfun Ġmark Reached Ops ( from Ops : ĠSet < Op >, Ġreached Ops : ĠMutable Set < Op >) Ġ{ ĊĠĠĠ Ġval Ġreached Queue Ġ= ĠArray Deque < Op >( from Ops ) ĊĠĠĠ Ġwhile Ġ( re ached Queue . isNotEmpty ()) Ġ{ ĊĠĠĠĠĠ Ġval Ġop Ġ= Ġreached Queue . remove First () ĊĠĠĠĠĠ Ġif Ġ( op Ġ! in Ġreached Ops ) Ġ{ ĊĠĠĠĠĠĠĠ Ġreached Ops Ġ+= Ġop ĊĠĠĠĠĠĠĠ Ġop . outputs . forEach Ġ{ Ġoutput Ġ-> ĊĠĠĠĠĠĠĠĠĠ Ġif Ġ( is Back propag atable ( output )) ĊĠĠĠĠĠĠĠĠĠĠĠ Ġreached Queue . addAll ( output . consum ers ) ĊĠĠĠĠĠĠĠ Ġ} ĊĠĠĠĠĠ Ġ} ĊĠĠĠ Ġ} ĊĠ Ġ} ĊĠĠĊĠ Ġprivate Ġval Ġtrain able Types Ġ= Ġset Of ( FLOAT 1 6 , ĠFLOAT , ĠDOUBLE , ĠCOMP LEX 6 4 , ĠCOMP LEX 1 2 8 , ĠRESOURCE ) ĊĠ Ġprivate Ġfun Ġis Train able ( tensor : ĠOutput ): ĠBoolean Ġ= ĊĠĠĠĠĠ Ġtensor . dataType . base DataType Ġin Ġtrain able Types ĊĠĠĊĠ Ġprivate Ġval Ġback propag atable Types Ġ= Ġset Of ( BF LOAT 1 6 , ĠRESOURCE , ĠVARI ANT ) ĊĠ Ġprivate Ġfun Ġis Back propag atable ( tensor : ĠOutput ): ĠBoolean Ġ= ĊĠĠĠĠĠ Ġis Train able ( tensor ) Ġ|| Ġtensor . dataType . base DataType Ġin Ġback propag atable Types ĊĠĠĊĠ Ġsealed Ġclass ĠG ating Method Ġ{ ĊĠĠĠ Ġobject ĠNo G ating Ġ: ĠG ating Method () ĊĠĠĠ Ġobject ĠOp G ating Ġ: ĠG ating Method () ĊĠĠĠ Ġobject ĠGraph G ating Ġ: ĠG ating Method () ĊĠ Ġ} ĊĠĠĊĠ Ġ/** ĊĠĠ Ġ* ĠAg gregation Ġmethod Ġused Ġto Ġcombine Ġgradients . ĊĠĠ Ġ* ĊĠĠ Ġ* ĠComputing Ġpartial Ġderivatives Ġcan Ġrequire Ġaggreg ating Ġgradient Ġcontributions . ĠAll Ġsuch Ġaggregation Ġmethods Ġare ĊĠĠ Ġ* Ġrepresented Ġas Ġobjects Ġextending Ġthis Ġtrait . ĊĠĠ Ġ* ĊĠĠ Ġ* Ġ@ see Ġ" tensorflow . python . ops . gradients _ impl . Aggregation Method " ĊĠĠ Ġ*/ ĊĠ Ġsealed Ġclass ĠAg gregation Method Ġ{ ĊĠĠĠĠĊĠĠĠ Ġ/** ĠAggregate Ġthe Ġgradients Ġfor Ġop Ġ` op `. ĊĠĠĠĠ Ġ* ĊĠĠĠĠ Ġ* Ġ@ param Ġ Ġgradients ĠĠ ĠMap Ġwhere Ġthe Ġcollected Ġgradients Ġare Ġstored . ĠThe Ġgradient Ġsequences Ġcorresponding Ġto Ġ` op ` ĊĠĠĠĠ Ġ* ĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġwill Ġbe Ġreplaced Ġwith Ġsequences Ġcontaining Ġa Ġsingle Ġelement Ġcorresponding Ġto Ġthe Ġaggregated ĊĠĠĠĠ Ġ* ĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġgradient . ĊĠĠĠĠ Ġ* Ġ@ param Ġ Ġop ĠĠĠĠĠĠĠĠĠ ĠOp Ġwhose Ġgradients Ġto Ġaggregate . ĊĠĠĠĠ Ġ* Ġ@ param Ġ Ġgradient UID ĠUnique Ġidentifier Ġwithin Ġthe Ġgraph Ġindicating Ġwhich Ġinvocation Ġof Ġgradients Ġis Ġbeing ĊĠĠĠĠ Ġ* ĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġexecuted . ĠUsed Ġto Ġcluster Ġops Ġfor Ġcompilation . ĊĠĠĠĠ Ġ* Ġ@ see Ġ" tensorflow . python . ops . gradients _ impl ._ Ag gregated Grad s " ĊĠĠĠĠ Ġ*/ ĊĠĠĠ Ġinternal Ġfun Ġaggregate Grad ients ( gradients : ĠMutable Map < Op , ĠMutable List < Mutable List < Output Like ? >> >, ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġop : ĠOp , ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġgradient UID : ĠString ): ĠMutable List < Mutable List < Output Like ? >> Ġ{ ĊĠĠĠĠĠ Ġval Ġop Grad ients Ġ= Ġgradients . get OrDefault ( op , Ġmutable ListOf ()) ĊĠĠĠĠĠ Ġreturn Ġif Ġ( control _ flow _ ops . is Loop Switch ( op )) ĊĠĠĠĠĠĠĠ Ġop Grad ients ĊĠĠĠĠĠ Ġelse Ġ{ ĊĠĠĠĠĠĠĠ Ġop Grad ients . with Index (). forEach Ġ{ Ġ( index , Ġgrad s ) Ġ-> ĊĠĠĠĠĠĠĠĠĠ Ġif Ġ( grad s . size Ġ< Ġ 2 ) ĊĠĠĠĠĠĠĠĠĠĠĠ Ġgrad s ĊĠĠĠĠĠĠĠĠĠ Ġelse ĊĠĠĠĠĠĠĠĠĠĠĠ Ġop Grad ients [ index ] Ġ= Ġmutable ListOf ( agg reate ( grad s . filter NotNull To ( mutable ListOf ()), Ġgradient UID )) ĊĠĠĠĠĠĠĠ Ġ} ĊĠĠĠĠĠĠĠ Ġop Grad ients ĊĠĠĠĠĠ Ġ} ĊĠĠĠ Ġ} ĊĠĠĠĠĊĠĠĠ Ġ/** ĠAg greg ates Ġ` values ` Ġinto Ġa Ġsingle Ġtensor . ĊĠĠĠĠ Ġ* ĊĠĠĠĠ Ġ* Ġ@ param Ġ Ġvalues ĠĠĠĠĠ ĠSequence Ġof Ġvalues Ġto Ġaggregate . ĊĠĠĠĠ Ġ* Ġ@ param Ġ Ġgradient UID ĠUnique Ġidentifier Ġwithin Ġthe Ġgraph Ġindicating Ġwhich Ġinvocation Ġof Ġgradients Ġis Ġbeing ĊĠĠĠĠ Ġ* ĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġexecuted Ġ( if Ġany ). ĠUsed Ġto Ġcluster Ġops Ġfor Ġcompilation . ĊĠĠĠĠ Ġ* Ġ@ return ĠAg gregated Ġtensor . ĊĠĠĠĠ Ġ*/ ĊĠĠĠ Ġinternal Ġabstract Ġfun Ġagg reate ( gradients : ĠMutable List < Output Like >, ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġgradient UID : ĠString ? Ġ= Ġnull ): ĠOutput Like ĊĠĠĠĠĊĠĠĠ Ġ/** ĠGradient Ġaggregation Ġmethod Ġthat Ġsimply Ġadds Ġup Ġthe Ġcollected Ġgradients . Ġ*/ ĊĠĠĠ Ġobject ĠAdd Aggregation Method Ġ: ĠAg gregation Method () Ġ{ ĊĠĠĠĠĠĠĊĠĠĠĠĠ Ġoverride Ġfun Ġagg reate ( gradients : ĠMutable List < Output Like >, Ġgradient UID : ĠString ?): ĠOutput Like Ġ{ ĊĠĠĠĠĠĠĠ Ġreturn Ġwhen Ġ{ ĊĠĠĠĠĠĠĠĠĠ Ġgradients . all Ġ{ Ġit Ġis ĠOutput Ġ} Ġ-> Ġ{ ĊĠĠĠĠĠĠĠĠĠĠĠ Ġ// ĠThis Ġfunction Ġadds Ġop Ġoutputs Ġfrom Ġpotentially Ġdifferent Ġdevices . ĊĠĠĠĠĠĠĠĠĠĠĠ Ġ// ĠWe Ġadd Ġthe Ġtensors Ġof Ġeach Ġdevice Ġseparately Ġfirst , Ġand Ġwe Ġthen Ġadd Ġup Ġthe Ġpartial Ġresults . ĊĠĠĠĠĠĠĠĠĠĠĠ Ġval Ġdevice Cont ributions Ġ= Ġgradients ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ. groupBy Ġ{ Ġit . device Ġ} ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ. toList () ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ. sorted By Ġ{ Ġit . first Ġ} ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ. map Ġ{ Ġ(_, Ġoutputs ) Ġ-> ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġtf . co locate With For Gradient ( mutable Set Of ( gradients [ 0 ]. op ), Ġgradient UID , Ġignore Existing Ġ= Ġtrue ) Ġ{ ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġtf . add N ( outputs . map Ġ{ Ġit Ġas ĠOutput Ġ}) ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ} ĊĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ} ĊĠĠĠĠĠĠĠĠĠĠĠ Ġtf . add N ( device Cont ributions ) ĊĠĠĠĠĠĠĠĠĠ Ġ} ĊĠĠĠĠĠĠĠĠĠ Ġgradients . all Ġ{ Ġit Ġis ĠIndexed S lices Ġ} Ġ-> Ġ{ ĊĠĠĠĠĠĠĠĠĠĠĠ ĠTODO () ĊĠĠĠĠĠĠĠĠĠ Ġ} ĊĠĠĠĠĠĠĠĠĠ Ġelse Ġ-> Ġthrow ĠIllegalArgumentException ( ĊĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ" The Ġgradients Ġbeing Ġaggregated Ġneed Ġto Ġbe Ġall Ġof Ġtype Ġ' Output ' Ġor Ġ' Indexed S lices '. ") ĊĠĠĠĠĠĠĠ Ġ} ĊĠĠĠĠĠ Ġ} ĊĠĠĠ Ġ} ĊĠĠĠĠĊĠĠĠ Ġ/** ĠGradient Ġaggregation Ġmethod Ġthat Ġsimply Ġadds Ġup Ġthe Ġcollected Ġgradients , Ġwithout Ġfirst Ġwaiting Ġfor Ġall Ġof Ġthem Ġto ĊĠĠĠĠ Ġ* Ġbecome Ġavailable Ġat Ġonce . ĊĠĠĠĠ Ġ* ĊĠĠĠĠ Ġ* ĠThe Ġbenefit Ġof Ġusing Ġthis Ġmethod Ġis Ġthat Ġits Ġinputs Ġcan Ġbe Ġcombined Ġin Ġany Ġorder Ġand Ġthis Ġcan Ġallow Ġthe Ġexpression ĊĠĠĠĠ Ġ* Ġto Ġbe Ġevaluated Ġwith Ġa Ġsmaller Ġmemory Ġfoot print . ĠWith Ġthis Ġmethod , Ġit Ġis Ġpossible Ġto Ġcompute Ġa Ġsum Ġof Ġterms Ġwhich ĊĠĠĠĠ Ġ* Ġare Ġmuch Ġlarger Ġthan Ġtotal ĠGPU Ġmemory . ĊĠĠĠĠ Ġ*/ ĊĠĠĠ Ġobject ĠAcc um ulate Aggregation Method Ġ: ĠAg gregation Method () Ġ{ ĊĠĠĠĠĠĠĊĠĠĠĠĠ Ġoverride Ġfun Ġagg reate ( gradients : ĠMutable List < Output Like >, Ġgradient UID : ĠString ?): ĠOutput Like Ġ{ ĊĠĠĠĠĠĠĠ Ġreturn Ġwhen Ġ{ ĊĠĠĠĠĠĠĠĠĠ Ġgradients . all Ġ{ Ġit Ġis ĠOutput Ġ} Ġ-> Ġtf . accum ulate N ( gradients . map Ġ{ Ġit Ġas ĠOutput Ġ}) ĊĠĠĠĠĠĠĠĠĠ Ġgradients . all Ġ{ Ġit Ġis ĠIndexed S lices Ġ} Ġ-> Ġ{ ĊĠĠĠĠĠĠĠĠĠĠĠ ĠTODO () ĊĠĠĠĠĠĠĠĠĠ Ġ} ĊĠĠĠĠĠĠĠĠĠ Ġelse Ġ-> Ġthrow ĠIllegalArgumentException ( ĊĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ" The Ġgradients Ġbeing Ġaggregated Ġneed Ġto Ġbe Ġall Ġof Ġtype Ġ' Output ' Ġor Ġ' Indexed S lices '. ") ĊĠĠĠĠĠĠĠ Ġ} ĊĠĠĠĠĠ Ġ} ĊĠĠĠĠĠĠĊĠĠĠ Ġ} ĊĠĠĠĠĊĠ Ġ} ĊĠĠĊĠ Ġ/** ĠRegistry Ġthat Ġcontains Ġthe Ġgradient Ġfunctions Ġto Ġbe Ġused Ġwhen Ġcreating Ġgradient Ġops . ĠGradient Ġfunctions Ġfor Ġall ĊĠĠ Ġ* Ġtypes Ġof Ġops Ġthat Ġare Ġbeing Ġdifferent i ated Ġneed Ġto Ġbe Ġregistered Ġusing Ġeither Ġthe Ġ[[ Registry . register ]] Ġor Ġthe ĊĠĠ Ġ* Ġ[[ Registry . register Non Different iable ]] Ġfunctions . ĠIn Ġan Ġattempt Ġto Ġobtain Ġthe Ġgradient Ġof Ġan Ġop Ġwhose Ġtype Ġhas Ġno ĊĠĠ Ġ* Ġgradient Ġfunction Ġregistered , Ġa Ġ[[ NoSuch ElementException ]] Ġwill Ġbe Ġthrown . Ġ*/ ĊĠ Ġobject ĠRegistry Ġ{ ĊĠĠĠĠĊĠĠĠ Ġprivate Ġval Ġregistry Ġ= Ġmutable Map Of < String , ĠGradient Function ?> () ĊĠĠĠĠĊĠĠĠ Ġ/** ĠRegisters Ġthe Ġprovided Ġgradient Ġfunction Ġto Ġthe Ġgradient Ġfunction Ġregistry . ĊĠĠĠĠ Ġ* ĊĠĠĠĠ Ġ* ĠNote Ġthat Ġif Ġa Ġgradient Ġfunction Ġfor Ġan Ġop Ġof Ġthe Ġsame Ġtype Ġalready Ġexists Ġin Ġthe Ġregistry , Ġthen Ġit Ġwill Ġbe ĊĠĠĠĠ Ġ* Ġover rid en Ġby Ġthe Ġprovided Ġgradient Ġfunction . ĊĠĠĠĠ Ġ* ĊĠĠĠĠ Ġ* Ġ@ param Ġ Ġop Type ĠĠ ĠOp Ġtype Ġfor Ġwhich Ġa Ġgradient Ġfunction Ġis Ġbeing Ġregistered . ĊĠĠĠĠ Ġ* Ġ@ param Ġ Ġfunction ĠGradient Ġfunction Ġ( t akes Ġop Ġand Ġoutput Ġgradients Ġas Ġinputs Ġand Ġreturns Ġthe Ġinput Ġgradients ). ĊĠĠĠĠ Ġ* Ġ@ see Ġ" tensorflow . python . framework . ops . Register Gradient " ĊĠĠĠĠ Ġ*/ ĊĠĠĠ Ġfun Ġregister ( var arg Ġop Types : ĠString , Ġfunction : ĠGradient Function ) Ġ{ ĊĠĠĠĠĠ Ġfor Ġ( op Type Ġin Ġop Types ) ĊĠĠĠĠĠĠĠ Ġregistry [ op Type ] Ġ= Ġfunction ĊĠĠĠ Ġ} ĊĠĠĠĠĊĠĠĠ Ġ/** ĠRegisters Ġthe Ġprovided Ġop Ġtype Ġas Ġnon - different iable Ġ( i . e ., Ġhaving Ġ` null ` Ġas Ġits Ġregistered Ġgradient Ġfunction ). ĊĠĠĠĠ Ġ* ĊĠĠĠĠ Ġ* ĠThis Ġfunction Ġshould Ġ* not * Ġbe Ġused Ġfor Ġops Ġthat Ġhave Ġa Ġwell - defined Ġgradient Ġthat Ġis Ġnot Ġyet Ġimplemented . ĠIt ĊĠĠĠĠ Ġ* Ġshould Ġonly Ġbe Ġused Ġwhen Ġdefining Ġa Ġnew Ġop Ġtype . ĠIt Ġmay Ġbe Ġused Ġfor Ġops Ġsuch Ġas Ġ` size ` Ġthat Ġare Ġnot ĊĠĠĠĠ Ġ* Ġdifferent iable . ĊĠĠĠĠ Ġ* ĊĠĠĠĠ Ġ* ĠThe Ġgradient Ġcomputed Ġfor Ġ' op Type ' Ġwill Ġthen Ġpropagate Ġzeros . ĊĠĠĠĠ Ġ* ĊĠĠĠĠ Ġ* ĠFor Ġops Ġthat Ġhave Ġa Ġwell - defined Ġgradient Ġbut Ġare Ġnot Ġyet Ġimplemented , Ġno Ġdeclaration Ġshould Ġbe Ġmade , Ġand Ġan Ġerror ĊĠĠĠĠ Ġ* Ġ* must * Ġbe Ġthrown Ġif Ġan Ġattempt Ġto Ġrequest Ġtheir Ġgradient Ġis Ġmade . ĊĠĠĠĠ Ġ* ĊĠĠĠĠ Ġ* Ġ@ param Ġ Ġop Type ĠOp Ġtype Ġto Ġregister Ġas Ġnon - different iable . ĊĠĠĠĠ Ġ*/ ĊĠĠĠ Ġfun Ġregister Non Different iable ( var arg Ġop Types : ĠString ) Ġ{ ĊĠĠĠĠĠ Ġfor Ġ( op Type Ġin Ġop Types ) ĊĠĠĠĠĠĠĠ Ġregistry [ op Type ] Ġ= Ġnull ĊĠĠĠ Ġ} ĊĠĠĠĠĊĠĠĠ Ġ/** ĠGets Ġthe Ġregistered Ġgradient Ġfunction Ġfor Ġthe Ġprovided Ġop Ġtype . ĊĠĠĠĠ Ġ* ĊĠĠĠĠ Ġ* Ġ@ param Ġ Ġop Type ĠOp Ġtype Ġwhose Ġgradient Ġfunction Ġis Ġbeing Ġlooked Ġup . ĊĠĠĠĠ Ġ* Ġ@ return ĠGradient Ġfunction Ġregistered Ġfor Ġthe Ġprovided Ġop Ġtype . ĊĠĠĠĠ Ġ* Ġ@ throws ĠNoSuch ElementException ĠIf Ġno Ġgradient Ġhas Ġbeen Ġregistered Ġfor Ġthe Ġprovided Ġop Ġtype . ĊĠĠĠĠ Ġ*/ ĊĠĠĠ Ġoperator Ġfun Ġget ( op Type : ĠString ): ĠGradient Function ? Ġ{ ĊĠĠĠĠĠ Ġif Ġ( op Type Ġ! in Ġregistry ) ĊĠĠĠĠĠĠĠ Ġthrow ĠNoSuch ElementException (" No Ġgradient Ġregistered Ġfor Ġop Ġtype Ġ'$ op Type '. ") ĊĠĠĠĠĠ Ġreturn Ġregistry [ op Type ] ĊĠĠĠ Ġ} ĊĠ Ġ} Ċ } Ċ type alias ĠGradient Function Ġ= Ġ( Op , ĠList < Output Like ?> ) Ġ-> ĠList < Output Like ?> Ċ Ċ // fun ĠTF . gradients ( y : ĠOutput , Ġxs : ĠCollection < Output >): ĠList < Output > Ġ{ Ċ // Ġ Ġval Ġ_ xs Ġ= ĠTF _ Output ( xs . size . to Long ()) Ċ // Ġ Ġfor Ġ(( i , Ġx ) Ġin Ġxs . with Index ()) Ċ // ĠĠĠ Ġ_ xs . position ( i . to Long ()). oper ( x . op !!. c _ op ). index ( x . value Index ) Ċ // Ġ Ġval Ġdy Ġ= ĠTF _ Output ( xs . size . to Long ()) Ċ // Ġ Ġval Ġstatus Ġ= Ġnew Status () Ċ // Ċ // Ġ ĠTF _ Add Grad ients ( g . c _ graph , Ċ // ĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġy . as TF _ Output (), Ġ 1 , Ċ // ĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ_ xs . position ( 0 L ), Ġxs . size , Ċ // ĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġnull , Ġstatus , Ċ // ĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġdy ) Ċ // Ġ Ġthrow Exception IfNot Ok ( status ) Ċ // Ġ Ġreturn ĠMutable List ( xs . size ) Ġ{ Ċ // ĠĠĠ Ġval Ġoutput Ġ= Ġdy . position ( it . to Long ()) Ċ // ĠĠĠ Ġval Ġout _ type Ġ= ĠTF _ Operation OutputType ( output ) Ċ // ĠĠĠ ĠOutput ( Op ( g , Ġoutput . oper ()), Ġoutput . index ()) Ċ // Ġ Ġ} Ċ //} Ċ Ċ // fun ĠTF . gradients ( ys : ĠList < Output >, Ġxs : ĠCollection < Output >): ĠList < Output > Ġ{ Ċ // Ġ Ġval Ġ_ ys Ġ= ĠTF _ Output ( ys . size . to Long ()) Ċ // Ġ Ġfor Ġ(( i , Ġy ) Ġin Ġys . with Index ()) Ċ // ĠĠĠ Ġ_ ys . position ( i . to Long ()). oper ( y . op !!. c _ op ). index ( y . value Index ) Ċ // Ġ Ġval Ġ_ xs Ġ= ĠTF _ Output ( xs . size . to Long ()) Ċ // Ġ Ġfor Ġ(( i , Ġx ) Ġin Ġxs . with Index ()) Ċ // ĠĠĠ Ġ_ xs . position ( i . to Long ()). oper ( x . op !!. c _ op ). index ( x . value Index ) Ċ // Ġ Ġval Ġdy Ġ= ĠTF _ Output ( xs . size . to Long ()) Ċ // Ġ Ġval Ġstatus Ġ= Ġnew Status () Ċ // Ġ ĠTF _ Add Grad ients ( g . c _ graph , Ċ // ĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ_ ys , Ġ 1 , Ċ // ĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġ_ xs . position ( 0 L ), Ġxs . size , Ċ // ĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġnull , Ġstatus , Ċ // ĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġdy ) Ċ // Ġ Ġthrow Exception IfNot Ok ( status ) Ċ // Ġ Ġreturn ĠMutable List ( xs . size ) Ġ{ Ċ // ĠĠĠ Ġval Ġoutput Ġ= Ġdy . position ( it . to Long ()) Ċ // ĠĠĠ ĠOutput ( Op ( g , Ġoutput . oper ()), Ġoutput . index ()) Ċ // Ġ Ġ} Ċ //} Ċ Ċ // fun ĠTF . gradient Des cent Optimizer ( learning Rate : ĠFloat , Ċ // ĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġloss : ĠOutput , Ċ // ĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠĠ Ġname : ĠString Ġ= Ġ" Gradient Des cent "): ĠOp Ġ{ Ċ // Ġ Ġname Scope ( name ) Ġ{ Ċ // ĠĠĠ Ġval Ġdy Ġ= Ġgradients ( loss , Ġtrain ables ) Ċ // ĠĠĠ Ġval Ġalpha Ġ= Ġconst ( learning Rate , Ġ" learning _ rate ") Ċ // ĠĠĠ Ġval Ġapply Gradient Ġ= Ġmutable ListOf < Op >() Ċ // ĠĠĠ Ġfor Ġ(( i , Ġtrain able ) Ġin Ġtrain ables . with Index ()) Ċ // ĠĠĠĠĠ Ġapply Gradient Ġ+= Ġapply Gradient Des cent ( train able , Ġalpha , Ġdy [ i ]) Ċ // ĠĠĠ Ġreturn Ġno Op Dep ( apply Gradient , Ġscope Name For Once ()) Ċ // Ġ Ġ} Ċ //}