Variables on TPU have a few restrictions.
(getter, name, *args, **kwargs)
| 972 | saved_custom_getter = vscope.custom_getter |
| 973 | |
| 974 | def custom_getter(getter, name, *args, **kwargs): |
| 975 | """Variables on TPU have a few restrictions.""" |
| 976 | partitioner = kwargs["partitioner"] |
| 977 | if partitioner is not None: |
| 978 | kwargs["partitioner"] = None |
| 979 | logging.warning( |
| 980 | "Partitioned variables are not supported on TPU. Got " |
| 981 | "`partitioner` that is {} for variable {}. " |
| 982 | "Setting `partitioner` to `None`." |
| 983 | .format(partitioner, name)) |
| 984 | if saved_custom_getter is None: |
| 985 | return getter(name, *args, **kwargs) |
| 986 | else: |
| 987 | return saved_custom_getter(getter, name, *args, **kwargs) |
| 988 | |
| 989 | vscope.set_use_resource(True) |
| 990 | vscope.set_custom_getter(custom_getter) |
no test coverage detected