Patch¶

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.