tensorflow: AutoCastVariable.assign returns wrapped variable instead of casted version

System information

Have I written custom code (as opposed to using a stock example script provided in TensorFlow): yes
OS Platform and Distribution (e.g., Linux Ubuntu 16.04): macOS
TensorFlow installed from (source or binary): binary
TensorFlow version (use command below): 2.0.0, 2.1-nightly
Python version: 3.7

Describe the current behavior

AutoCastVariable variable forwards assign and scatter to the underlying float32 variable: https://github.com/tensorflow/tensorflow/blob/cee2a43b8184e92ba26ec0e3d6e00a3f8ca6e3c8/tensorflow/python/keras/mixed_precision/experimental/autocast_variable.py#L187-L188 Thus, the return value of assign methods with read_value=True is a normal tf.Variable and not an AutoCastVariable. This means that calculations directly depending on the assign operation, might run in float32 instead of float16, or am I missing something?

Describe the expected behavior AutoCastVariable.assign* should return an AutoCastVariable variable instead tf.Variable so that the dtype is preserved.

@reedwm Is this intended behaviour?

Code to reproduce the issue

import tensorflow as tf
from tensorflow.python.keras.mixed_precision.experimental import autocast_variable

var = tf.Variable(0., dtype=tf.float32)
var = autocast_variable.AutoCastVariable(var)

with tf.compat.v1.get_default_graph()._enable_auto_casting_variables(tf.float16):
    assert var.dtype == tf.float16
    # assign should return an AutoCastVariable but returns tf.Variable
    var_assign = var.assign(5.)
    assert not isinstance(var_assign, autocast_variable.AutoCastVariable)
    assert var_assign.dtype == tf.float32

About this issue

Original URL
State: closed
Created 5 years ago
Comments: 26 (24 by maintainers)

Most upvoted comments

This is still an issue. The tf.distribute.Strategy team is working on a way to return a variable from DistributionVariable.assign and until then, this issue cannot be fully fixed.

reedwm on Mar 16, 2020

This is a good idea, but I personally don’t think it’s worth the infrastructure complexity of running DISABLED tests to check whether a bug is fixed. Typically if a bug is fixed, someone will write a corresponding unit test. This bug is an exception, but even so, it was not fully fixed (it still doesn’t work with distribution strategies in graph mode).

reedwm on Jun 2, 2021

You’re right, reopening.

Once Distribution Strategy returns a variable from DistributionVariable.assign, this issue can be fixed.

reedwm on Dec 6, 2019

I could fix this by copying or subclassing _UnreadVariable. I could also just return self, which has slightly incorrect semantics when Sessions/Graphs are used with tf.compat.v1

@reedwm Since AutoCastVariable just wraps an other variable, what about returning a new instance wrapping the _UnreadVariable returned by the assign op. Would something like the following work?

def assign(self, value, use_locking=None, name=None, read_value=True):
    assign_op_or_var = self._variable.assign(value, use_locking, name, read_value)
    return self.__class__(assign_op_or_var) if read_value else assign_op_or_var
#   return AutoCastVariable(assign_op_or_var) if read_value else assign_op_or_var

lgeiger on Nov 29, 2019

I have tried on colab with TF version 2.0 ,2.1.0-dev20191111 and was able to reproduce the issue.Please, find the gist here. Thanks!

ravikyram on Nov 18, 2019