Patches for TensorFlow.

class PatchTensorFlow[source]

Bases: object

Patches for TensorFlow.

There are multiple “monkey patches”: 1) Patches the value attribute of a ResourceVariable to return the graph_element or cached_value instead of a ReadVariableOp. Not having to read the variable from memory again can provide performance increases. 2) Patches optimizers to save their type and arguments so that we can re-create them when partitioning variables. 3) Patches Keras’ Session fetcher to return AutoDist’s custom Session so that AutoDist is compatible with Keras.

static patch_var_reading()[source]

It only works with tf.gradients but not tape.gradients.

static unpatch_var_reading()[source]

Revert the ReadVariable patch.

static patch_optimizers()[source]

Patch all instances of OptimizerV2 for us to store optimizer and gradient information.

static patch_keras()[source]

Patch Keras way of getting session.

static unpatch_keras()[source]

Revert the Keras patch.