Home
last modified time | relevance | path

Searched refs:mutable_entry_computation_layout (Results 1 – 14 of 14) sorted by relevance

/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/tests/
H A Dhlo_test_base.h303 module->mutable_entry_computation_layout()->parameter_count()); in ForceParameterLayout()
304 module->mutable_entry_computation_layout() in ForceParameterLayout()
312 module->mutable_entry_computation_layout() in ForceResultLayout()
319 module->mutable_entry_computation_layout() in ForceResultLayout()
327 module->mutable_entry_computation_layout() in ForceClearResultLayout()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/
H A Dhlo_module_util.cc51 config->mutable_entry_computation_layout(); in CreateModuleConfig()
172 *module->mutable_entry_computation_layout()->mutable_parameter_layout(i) = in UpdateEntryComputationLayout()
177 *module->mutable_entry_computation_layout()->mutable_result_layout() = in UpdateEntryComputationLayout()
H A Dhlo_module.h148 ComputationLayout* mutable_entry_computation_layout() { in mutable_entry_computation_layout() function
149 return config_.mutable_entry_computation_layout(); in mutable_entry_computation_layout()
H A Dhlo_module_config.h106 ComputationLayout* mutable_entry_computation_layout() { in mutable_entry_computation_layout() function
H A Dlayout_assignment_test.cc1336 m->mutable_entry_computation_layout() in AssignLayoutsToComputation()
1340 LayoutAssignment layout_assignment(m->mutable_entry_computation_layout(), in AssignLayoutsToComputation()
1369 *m->mutable_entry_computation_layout()->mutable_result_layout() = in TEST_F()
H A Dhlo_computation.cc171 config.mutable_entry_computation_layout()->add_parameter_layout( in AddEntryComputationParameter()
191 *config.mutable_entry_computation_layout()->mutable_parameter_layout( in ReplaceEntryComputationParameter()
H A Dhlo_module.cc572 module_config.mutable_entry_computation_layout(); in CreateModuleConfigFromShape()
H A Ddynamic_padder_test.cc539 new_config.mutable_entry_computation_layout() in PadAndExecute()
H A Dhlo_parser.cc954 *config.mutable_entry_computation_layout() = *entry_computation_layout; in ParseHloModule()
1001 TF_CHECK_OK(module->mutable_entry_computation_layout() in ParseComputations()
1006 TF_CHECK_OK(module->mutable_entry_computation_layout() in ParseComputations()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/interpreter/
H A Dcompiler.cc96 hlo_module->mutable_entry_computation_layout()); in RunHloOptimization()
/aosp_15_r20/external/tensorflow/tensorflow/core/tpu/kernels/
H A Dtpu_compile_op_support.cc68 config->mutable_entry_computation_layout(); in CreateModuleConfig()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/cpu/
H A Dcpu_compiler.cc613 module->mutable_entry_computation_layout(), target_machine_features, in RunHloPassesThroughLayoutAssn()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/gpu/
H A Dgpu_compiler.cc599 hlo_module->mutable_entry_computation_layout(), stream_exec, in OptimizeHloModule()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/spmd/
H A Dspmd_partitioner.cc4556 *config.mutable_entry_computation_layout() = in Run()