def compile(self, c_computation, compile_options):
options = _xla.ExecutableBuildOptions()
options.num_replicas = compile_options.num_replicas
if compile_options.result_layout:
options.result_layout = compile_options.result_layout
options.debug_options.xla_cpu_fast_math_honor_infs = True
options.debug_options.xla_cpu_fast_math_honor_nans = True
options.debug_options.xla_cpu_fast_math_honor_division = True
options.debug_options.xla_cpu_fast_math_honor_functions = True
options.debug_options.xla_gpu_enable_fast_min_max = False
return _tpu_client.TpuExecutable.Compile(
c_computation, compile_options.argument_layouts, options,
self.client, compile_options.device_assignment)
def __init__(self):
self.executable_build_options = _xla.ExecutableBuildOptions()
self.xla_dump_to = None
self.dump_hlo_pass_re = None
self.dump_hlo_module_re = None
self.dump_hlo_as_text = None
self.dump_hlo_as_proto = None
self.hlo_profile = None
self.num_replicas = 1
self.num_partitions = 1
self.argument_layouts = None
self.result_layout = None
self.device_assignment = None
self.tuple_arguments = False
def compile(self, c_computation, compile_options=None):
compile_options = compile_options or CompileOptions()
options = _xla.ExecutableBuildOptions()
options.num_replicas = compile_options.num_replicas
options.num_partitions = compile_options.num_partitions
if compile_options.result_layout:
options.result_layout = compile_options.result_layout
options.debug_options.xla_cpu_fast_math_honor_infs = True
options.debug_options.xla_cpu_fast_math_honor_nans = True
options.debug_options.xla_cpu_fast_math_honor_division = True
options.debug_options.xla_cpu_fast_math_honor_functions = True
options.debug_options.xla_gpu_enable_fast_min_max = False
return _xla.LocalExecutable.compile(c_computation,
compile_options.argument_layouts,
options, self.client,
compile_options.device_assignment,
compile_options.tuple_arguments)
