lingvo.core.var_tmp_wrappers module

Wrappers of tf.Variable to achieve various temporary effects.

E.g., such effects include tracking assign ops, and converting between manual and auto sharding modes in GShard. These wrappers are intended to be used in limited scopes.

class lingvo.core.var_tmp_wrappers.VarWrapperTrackAssign(var)

Bases: object

A wrapper of tf.Variable that tracks assignments.

previous_assigns()

control_after_assigns()

property raw_var

assign(value, use_locking=False, name=None, read_value=True)

assign_add(delta, use_locking=False, name=None, read_value=True)

assign_sub(delta, use_locking=False, name=None, read_value=True)

class lingvo.core.var_tmp_wrappers.StackedVarWrapperWithManualSharding(var)

Bases: object

A wrapper of tf.Variable for stacked variables in manual-sharding mode.

The variable is sharded on on the leading (stacking) dimension, and the shard size is 1. Example: the physical variable v has shape [N, a, b], which is stacked from N logical variables of shape [a, b] and annotated to be sharded on dim 0 in N ways:

- With StackedVarWrapperWithManualSharding(v), read from the wrapper will
have (manually sharded) shape [a, b].

This wrapper internally converts between auto and manual sharding modes, which makes variable read/write compatible with the rest of manually sharded code.

If the variable has other dimensions sharded, the auto/manual conversion will only happen partially on the stacking dimension only.

property raw_var

_to_manual(val)

_to_auto(val)

value()

Returns the variable and converts it to manually sharded mode.

Returns

The return value has the shape of the individual elements of the stacked variable (shard shape with the stacking dimension collapsed).

read_value()

Reads the variable and converts it to manually sharded mode.

Returns

The return value has the shape of the individual elements of the stacked variable (shard shape with the stacking dimension collapsed).

assign(value, use_locking=False, name=None, read_value=True)

Implements the interface of tf.Variable.assign.

Parameters

• value – A manually sharded tensor that has the shape of the individual elements of the stacked variable (shard shape with the stacking dimension collapsed).

• use_locking – See tf.Variable.assign.

• name – See tf.Variable.assign.

• read_value – See tf.Variable.assign. If True, the returned value will be manually sharded.

Returns

See tf.Variable.assign. If read_value is True, returns the updated value in the shard shape of the shape of the individual elements of the stacked variable (shard shape with the stacking dimension collapsed).

assign_add(delta, use_locking=False, name=None, read_value=True)

Implements the interface of tf.Variable.assign_add.

assign_sub(delta, use_locking=False, name=None, read_value=True)

Implements the interface of tf.Variable.assign_sub.