(self, *args, **kwargs)
| 1002 | self._t.start() |
| 1003 | |
| 1004 | def result(self, timeout: Optional[int] = None) -> Any: |
| 1005 | return self._t.join(timeout=timeout) |
| 1006 | |
| 1007 | |
| 1008 | def _get_premapped_buffer_size(): |
| 1009 | if jax.default_backend() == "tpu": |
| 1010 | # If TPU_PREMAPPED_BUFFER_SIZE is not set, default is 4GB. |
| 1011 | return int(os.getenv("TPU_PREMAPPED_BUFFER_SIZE", "4294967296")) |
| 1012 | # On all other backends, use 1TB (effectively unlimited). |
| 1013 | return 1099511627776 |
| 1014 |