Commitcccc00e

chuanhaozhuge

authored and

Mesh TensorFlow Team

committed

Make TPU variable name deterministic.

Only use the unique variable handle id as handle cache key, but not the variable handle's name.This resolves a graph non-determinism issue seen in an internal workload.PiperOrigin-RevId: 436262849

1 parent39f4bd6 commitcccc00eCopy full SHA for cccc00e

File tree

1 file changed

-1

lines changed

mesh_tensorflow
- tpu_variables.py

1 file changed

-1

lines changed

`‎mesh_tensorflow/tpu_variables.py‎`

Lines changed: 3 additions & 1 deletion

Original file line number	Diff line number	Diff line change
`@@ -82,7 +82,9 @@ def handle(self):`
`82`	`82`	`iftpu_contextisNone:`
`83`	`83`	`returnself._primary_var.handle`
`84`	`84`
`85`		`-returntpu_context.get_replicated_var_handle(self._name,self._vars)`
	`85`	`+# Using variable name as handle id.`
	`86`	`+returntpu_context.get_replicated_var_handle(self._name,self._name,`
	`87`	`+self._vars)`
`86`	`88`
`87`	`89`	`@contextlib.contextmanager`
`88`	`90`	`def_assign_dependencies(self):`

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!

Commitcccc00e

File tree

1 file changed

1 file changed

`‎mesh_tensorflow/tpu_variables.py‎`

0 commit comments