Custom getter that forces variables to have type self.variable_type.
(getter, *args, **kwargs)
| 47 | """ |
| 48 | |
| 49 | def inner_custom_getter(getter, *args, **kwargs): |
| 50 | """Custom getter that forces variables to have type self.variable_type.""" |
| 51 | cast_to_bfloat16 = False |
| 52 | requested_dtype = kwargs['dtype'] |
| 53 | if requested_dtype == dtypes.bfloat16: |
| 54 | # Only change the variable dtype if doing so does not decrease variable |
| 55 | # precision. |
| 56 | kwargs['dtype'] = dtypes.float32 |
| 57 | cast_to_bfloat16 = True |
| 58 | var = getter(*args, **kwargs) |
| 59 | # This if statement is needed to guard the cast, because batch norm |
| 60 | # assigns directly to the return value of this custom getter. The cast |
| 61 | # makes the return value not a variable so it cannot be assigned. Batch |
| 62 | # norm variables are always in fp32 so this if statement is never |
| 63 | # triggered for them. |
| 64 | if cast_to_bfloat16: |
| 65 | var = math_ops.cast(var, dtypes.bfloat16) |
| 66 | return var |
| 67 | |
| 68 | return inner_custom_getter |
| 69 |