| 142 | return _xla.PyLocalBuffer.make_tuple(c_buffers, self.client, device) |
| 143 | |
| 144 | def compile(self, c_computation, compile_options): |
| 145 | options = _xla.ExecutableBuildOptions() |
| 146 | options.num_replicas = compile_options.num_replicas |
| 147 | options.num_partitions = compile_options.num_partitions |
| 148 | if compile_options.result_layout: |
| 149 | options.result_layout = compile_options.result_layout |
| 150 | options.debug_options.xla_cpu_fast_math_honor_infs = True |
| 151 | options.debug_options.xla_cpu_fast_math_honor_nans = True |
| 152 | options.debug_options.xla_cpu_fast_math_honor_division = True |
| 153 | options.debug_options.xla_cpu_fast_math_honor_functions = True |
| 154 | options.debug_options.xla_gpu_enable_fast_min_max = False |
| 155 | return _xla.LocalExecutable.Compile(c_computation, |
| 156 | compile_options.argument_layouts, |
| 157 | options, self.client, |
| 158 | compile_options.device_assignment) |
| 159 | |
| 160 | def get_default_device_assignment(self, num_replicas, num_partitions=None): |
| 161 | if num_partitions is not None: |