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)