(self, *args)
| 270 | self.kept_var_idx = kept_var_idx |
| 271 | |
| 272 | def __call__(self, *args): |
| 273 | args = [x for i, x in enumerate(args) if i in self.kept_var_idx] |
| 274 | input_bufs = self.in_handler(args) |
| 275 | assert not ( |
| 276 | self.ordered_effects |
| 277 | or self.has_unordered_effects |
| 278 | or self.has_host_callbacks |
| 279 | ) |
| 280 | |
| 281 | if True or not is_distributed(): |
| 282 | out_bufs = self.xla_executable.execute_sharded_on_local_devices(input_bufs) |
| 283 | return self.out_handler(out_bufs) |
| 284 | else: |
| 285 | results = self.xla_executable.execute_sharded(input_bufs) |
| 286 | outputs = results.disassemble_into_single_device_arrays() |
| 287 | assert isinstance(outputs, list) |
| 288 | out_bufs = [] |
| 289 | for oup in outputs: |
| 290 | assert isinstance(oup, list) and len(oup) == 1 |
| 291 | out_bufs.append(oup[0].device_buffers) |
| 292 | return self.out_handler(out_bufs) |
| 293 | |
| 294 | |
| 295 | @dataclasses.dataclass |
nothing calls this directly
no test coverage detected