Commit6d44045

authored and

Mesh TensorFlow Team

committed

Make mesh_tensorflow's call ofget_replicated_var_handle backward-compatible with tf <= 2.8.0.Fixesgoogle-research/text-to-text-transfer-transformer#1020.

PiperOrigin-RevId: 448301001

1 parentb631ad5 commit6d44045Copy full SHA for 6d44045

File tree

-4

lines changed

-4

lines changed

Lines changed: 8 additions & 3 deletions

Original file line number	Diff line number	Diff line change
`@@ -20,6 +20,7 @@`
`20`	`20`	`from __future__importprint_function`
`21`	`21`
`22`	`22`	`importcontextlib`
	`23`	`+importinspect`
`23`	`24`
`24`	`25`	`# pylint: disable=g-direct-tensorflow-import`
`25`	`26`	`fromtensorflow.python.frameworkimportops`
`@@ -82,9 +83,13 @@ def handle(self):`
`82`	`83`	`iftpu_contextisNone:`
`83`	`84`	`returnself._primary_var.handle`
`84`	`85`
`85`		`-# Using variable name as handle id.`
`86`		`-returntpu_context.get_replicated_var_handle(self._name,self._name,`
`87`		`-self._vars)`
	`86`	`+# TODO(adarob): Remove backward-compatibility when TF 2.10 is released.`
	`87`	`+if'handle_id'notininspect.signature(`
	`88`	`+tpu_context.get_replicated_var_handle).parameters:`
	`89`	`+returntpu_context.get_replicated_var_handle(`
	`90`	`+name=self._name,vars_=self._vars)`
	`91`	`+returntpu_context.get_replicated_var_handle(`
	`92`	`+name=self._name,handle_id=self._name,vars_=self._vars)`
`88`	`93`
`89`	`94`	`@contextlib.contextmanager`
`90`	`95`	`def_assign_dependencies(self):`

Lines changed: 1 addition & 1 deletion

Original file line number	Diff line number	Diff line change
`@@ -5,7 +5,7 @@`
`5`	`5`
`6`	`6`	`setup(`
`7`	`7`	`name='mesh-tensorflow',`
`8`		`-version='0.1.20',`
	`8`	`+version='0.1.21',`
`9`	`9`	`description='Mesh TensorFlow',`
`10`	`10`	`author='Google Inc.',`
`11`	`11`	`author_email='no-reply@google.com',`

Comments

(0)