(self, getter, *args, **kwargs)
| 357 | self.small_variable_size_threshold = small_variable_size_threshold |
| 358 | |
| 359 | def __call__(self, getter, *args, **kwargs): |
| 360 | size = tf.TensorShape(kwargs['shape']).num_elements() |
| 361 | if size is None or not kwargs.get('trainable', True): |
| 362 | # TODO a lot of vars won't be saved then |
| 363 | _replace_global_by_local(kwargs) |
| 364 | return getter(*args, **kwargs) |
| 365 | |
| 366 | if size < self.small_variable_size_threshold: |
| 367 | device_name = self.device_for_small_variables |
| 368 | else: |
| 369 | device_index, _ = min(enumerate(self.sizes), key=operator.itemgetter(1)) |
| 370 | device_name = self.devices[device_index] |
| 371 | self.sizes[device_index] += size |
| 372 | |
| 373 | kwargs['caching_device'] = device_name |
| 374 | var = getter(*args, **kwargs) |
| 375 | return var |
| 376 | |
| 377 | |
| 378 | # TODO pack at variable boundary, so that the concat does not have to wait for all |
nothing calls this directly
no test coverage detected