Method keep_weights

tensorflow/python/ops/variable_scope.py:3008–3030

Scope class for bfloat16 variables so that the model uses custom getter. This enables variables to be read as bfloat16 type when using get_variable. ```python import tensorflow as tf from tensorflow.contrib import layers with tf.variable_scope(...).keep_weights(dtype=tf.float3

(self, dtype=dtypes.float32)

    return inner_custom_getter

  def keep_weights(self, dtype=dtypes.float32):
    """Scope class for bfloat16 variables so that the model uses custom getter.

    This enables variables to be read as bfloat16 type when using get_variable.

    ```python
    import tensorflow as tf
    from tensorflow.contrib import layers

    with tf.variable_scope(...).keep_weights(dtype=tf.float32):
      data_bf16 = tf.cast(data, dtype=tf.bfloat16)

      matmul_0 = tf.layers.dense(data_bf16, 64, activation=tf.nn.relu)
      matmul_0 = tf.layers.batch_normalization(matmul_0, training=True)
      matmul_0 = tf.cast(matmul_0, dtype=tf.float32)

      matmul_1 = layers.fully_connected(data_bf16, 128,
                                        activation_fn=tf.nn.leaky_relu)
      matmul_1 = tf.cast(matmul_1, dtype=tf.float32)
    ```
    """
    self._custom_getter = self._get_custom_getter(dtype=dtype)
    return self

  def __enter__(self):
    # If the default graph is building a function, then we should not replace it