NotificationsYou must be signed in to change notification settings
Fork2.2k
Star9.4k

Commit7a734f4

committed

Tweaks to model_graph to play nice with XTensorVariables

* Use RV Op name when provided* More robust detection of observed data variables (after#7656 arbitrary graphs are allowed)* Remove self loops explicitly (closes#7722)

1 parentc3cc833 commit7a734f4Copy full SHA for 7a734f4

File tree

1 file changed

+18

-16

lines changed

pymc
- model_graph.py

1 file changed

+18

-16

lines changed

`‎pymc/model_graph.py‎`

Lines changed: 18 additions & 16 deletions

Original file line number	Diff line number	Diff line change
`@@ -21,11 +21,11 @@`
`21`	`21`	`fromtypingimportAny,cast`
`22`	`22`
`23`	`23`	`frompytensorimportfunction`
`24`		`-frompytensor.graph.basicimportancestors,walk`
	`24`	`+frompytensor.graph.basicimportVariable,ancestors,walk`
`25`	`25`	`frompytensor.tensor.shapeimportShape`
`26`		`-frompytensor.tensor.variableimportTensorVariable`
`27`	`26`
`28`	`27`	`frompymc.model.coreimportmodelcontext`
	`28`	`+frompymc.pytensorfimport_cheap_eval_mode`
`29`	`29`	`frompymc.utilimportVarName,get_default_varnames,get_var_name`
`30`	`30`
`31`	`31`	`__all__= (`
`@@ -73,7 +73,7 @@ def create_plate_label_with_dim_length(`
`73`	`73`
`74`	`74`
`75`	`75`	`deffast_eval(var):`
`76`		`-returnfunction([],var,mode="FAST_COMPILE")()`
	`76`	`+returnfunction([],var,mode=_cheap_eval_mode)()`
`77`	`77`
`78`	`78`
`79`	`79`	`classNodeType(str,Enum):`
`@@ -88,7 +88,7 @@ class NodeType(str, Enum):`
`88`	`88`
`89`	`89`	`@dataclass`
`90`	`90`	`classNodeInfo:`
`91`		`-var:TensorVariable`
	`91`	`+var:Variable`
`92`	`92`	`node_type:NodeType`
`93`	`93`
`94`	`94`	`def__hash__(self):`
`@@ -108,10 +108,10 @@ def __eq__(self, other) -> bool:`
`108`	`108`
`109`	`109`
`110`	`110`	`GraphvizNodeKwargs=dict[str,Any]`
`111`		`-NodeFormatter=Callable[[TensorVariable],GraphvizNodeKwargs]`
	`111`	`+NodeFormatter=Callable[[Variable],GraphvizNodeKwargs]`
`112`	`112`
`113`	`113`
`114`		`-defdefault_potential(var:TensorVariable)->GraphvizNodeKwargs:`
	`114`	`+defdefault_potential(var:Variable)->GraphvizNodeKwargs:`
`115`	`115`	`"""Return default data for potential in the graph."""`
`116`	`116`	`return {`
`117`	`117`	`"shape":"octagon",`
`@@ -120,17 +120,19 @@ def default_potential(var: TensorVariable) -> GraphvizNodeKwargs:`
`120`	`120`	`}`
`121`	`121`
`122`	`122`
`123`		`-defrandom_variable_symbol(var:TensorVariable)->str:`
	`123`	`+defrandom_variable_symbol(var:Variable)->str:`
`124`	`124`	`"""Get the symbol of the random variable."""`
`125`		`-symbol=var.owner.op.__class__.__name__`
	`125`	`+op=var.owner.op`
`126`	`126`
`127`		`-ifsymbol.endswith("RV"):`
`128`		`-symbol=symbol[:-2]`
	`127`	`+ifname:=getattr(op,"name",None):`
	`128`	`+symbol=name[0].upper()+name[1:]`
	`129`	`+else:`
	`130`	`+symbol=op.__class__.__name__.removesuffix("RV")`
`129`	`131`
`130`	`132`	`returnsymbol`
`131`	`133`
`132`	`134`
`133`		`-defdefault_free_rv(var:TensorVariable)->GraphvizNodeKwargs:`
	`135`	`+defdefault_free_rv(var:Variable)->GraphvizNodeKwargs:`
`134`	`136`	`"""Return default data for free RV in the graph."""`
`135`	`137`	`symbol=random_variable_symbol(var)`
`136`	`138`
`@@ -141,7 +143,7 @@ def default_free_rv(var: TensorVariable) -> GraphvizNodeKwargs:`
`141`	`143`	`}`
`142`	`144`
`143`	`145`
`144`		`-defdefault_observed_rv(var:TensorVariable)->GraphvizNodeKwargs:`
	`146`	`+defdefault_observed_rv(var:Variable)->GraphvizNodeKwargs:`
`145`	`147`	`"""Return default data for observed RV in the graph."""`
`146`	`148`	`symbol=random_variable_symbol(var)`
`147`	`149`
`@@ -152,7 +154,7 @@ def default_observed_rv(var: TensorVariable) -> GraphvizNodeKwargs:`
`152`	`154`	`}`
`153`	`155`
`154`	`156`
`155`		`-defdefault_deterministic(var:TensorVariable)->GraphvizNodeKwargs:`
	`157`	`+defdefault_deterministic(var:Variable)->GraphvizNodeKwargs:`
`156`	`158`	`"""Return default data for the deterministic in the graph."""`
`157`	`159`	`return {`
`158`	`160`	`"shape":"box",`
`@@ -161,7 +163,7 @@ def default_deterministic(var: TensorVariable) -> GraphvizNodeKwargs:`
`161`	`163`	`}`
`162`	`164`
`163`	`165`
`164`		`-defdefault_data(var:TensorVariable)->GraphvizNodeKwargs:`
	`166`	`+defdefault_data(var:Variable)->GraphvizNodeKwargs:`
`165`	`167`	`"""Return default data for the data in the graph."""`
`166`	`168`	`return {`
`167`	`169`	`"shape":"box",`
`@@ -239,7 +241,7 @@ def __init__(self, model):`
`239`	`241`	`self._all_vars= {model[var_name]forvar_nameinself._all_var_names}`
`240`	`242`	`self.var_list=self.model.named_vars.values()`
`241`	`243`
`242`		`-defget_parent_names(self,var:TensorVariable)->set[VarName]:`
	`244`	`+defget_parent_names(self,var:Variable)->set[VarName]:`
`243`	`245`	`ifvar.ownerisNone:`
`244`	`246`	`returnset()`
`245`	`247`
`@@ -345,7 +347,7 @@ def get_plates(`
`345`	`347`	`dim_name:fast_eval(value).item()fordim_name,valueinself.model.dim_lengths.items()`
`346`	`348`	`}`
`347`	`349`	`var_shapes:dict[str,tuple[int, ...]]= {`
`348`		`-var_name:tuple(fast_eval(self.model[var_name].shape))`
	`350`	`+var_name:tuple(map(int,fast_eval(self.model[var_name].shape)))`
`349`	`351`	`forvar_nameinself.vars_to_plot(var_names)`
`350`	`352`	`}`
`351`	`353`

0 commit comments

Comments

(0)

Movatterモバイル変換

Navigation Menu

Search code, repositories, users, issues, pull requests...

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Commit7a734f4

File tree

1 file changed

1 file changed

`‎pymc/model_graph.py‎`

0 commit comments