Home
last modified time | relevance | path

Searched refs:program_shape (Results 1 – 25 of 31) sorted by relevance

12

/external/tensorflow/tensorflow/compiler/tf2xla/
Dxla_jit_compiled_cpu_function.cc40 const xla::ProgramShape& program_shape) { in ComputeArgSizes() argument
42 const size_t num_args = program_shape.parameters_size(); in ComputeArgSizes()
45 const xla::Shape& arg_shape = program_shape.parameters(i); in ComputeArgSizes()
123 TF_ASSIGN_OR_RETURN(std::unique_ptr<xla::ProgramShape> program_shape, in Compile()
125 if (program_shape->result().element_type() != xla::TUPLE) { in Compile()
134 program_shape->clear_parameter_names(); in Compile()
138 arg_shapes.reserve(program_shape->parameters_size()); in Compile()
139 for (int i = 0; i < program_shape->parameters_size(); ++i) { in Compile()
140 arg_shapes.push_back(&program_shape->parameters(i)); in Compile()
157 ComputeArgSizes(*program_shape)); in Compile()
[all …]
Dxla_jit_compiled_cpu_function_test.cc119 const xla::ProgramShape* program_shape = function.ProgramShape(); in TEST() local
120 ASSERT_TRUE(program_shape != nullptr); in TEST()
121 ASSERT_EQ(program_shape->parameters_size(), 2); in TEST()
122 EXPECT_TRUE(ShapeUtil::Compatible(program_shape->parameters(0), s32)); in TEST()
123 EXPECT_TRUE(ShapeUtil::Compatible(program_shape->parameters(1), s32)); in TEST()
125 const xla::Shape& result = program_shape->result(); in TEST()
Dxla_compiled_cpu_function.cc31 program_shape_(static_data.program_shape), in XlaCompiledCpuFunction()
/external/tensorflow/tensorflow/compiler/xla/service/
Dlocal_service.cc82 std::shared_ptr<const ProgramShape> program_shape, in CompileExecutable()
86 if (argument_layouts.size() != program_shape->parameters_size()) { in CompileExecutable()
89 program_shape->parameters_size(), argument_layouts.size()); in CompileExecutable()
94 if (!ShapeUtil::Compatible(argument_shape, program_shape->parameters(i))) { in CompileExecutable()
112 ShapeUtil::HumanString(program_shape->parameters(i)).c_str(), in CompileExecutable()
118 *build_options.result_layout(), program_shape->result())); in CompileExecutable()
131 program_shape->result(); in CompileExecutable()
137 CreateModuleConfig(*program_shape, argument_layouts, &execution_options, in CompileExecutable()
Dhlo_module_config.cc33 HloModuleConfig::HloModuleConfig(const ProgramShape& program_shape) in HloModuleConfig() argument
34 : entry_computation_layout_(program_shape) {} in HloModuleConfig()
37 const ProgramShape& program_shape) { in SetDefaultComputationLayout() argument
38 entry_computation_layout_ = ComputationLayout(program_shape); in SetDefaultComputationLayout()
Duser_computation_test.cc86 std::shared_ptr<const ProgramShape> program_shape, in TEST_F()
88 ASSERT_EQ(1, program_shape->parameters_size()); in TEST_F()
90 ShapeUtil::Compatible(kScalarShape, program_shape->parameters(0))); in TEST_F()
91 EXPECT_TRUE(ShapeUtil::Compatible(kScalarShape, program_shape->result())); in TEST_F()
113 std::shared_ptr<const ProgramShape> program_shape, in TEST_F()
115 ASSERT_EQ(1, program_shape->parameters_size()); in TEST_F()
117 ShapeUtil::Compatible(kScalarShape, program_shape->parameters(0))); in TEST_F()
118 EXPECT_TRUE(ShapeUtil::Compatible(kScalarShape, program_shape->result())); in TEST_F()
Dcomputation_layout.cc26 ComputationLayout::ComputationLayout(const ProgramShape& program_shape) in ComputationLayout() argument
27 : result_layout_(program_shape.result()) { in ComputationLayout()
28 for (auto& shape : program_shape.parameters()) { in ComputationLayout()
Dhlo_module.cc260 ProgramShape program_shape; in ProgramShapeFromProto() local
261 *program_shape.mutable_result() = root->shape(); in ProgramShapeFromProto()
267 *program_shape.add_parameters() = shape; in ProgramShapeFromProto()
268 program_shape.add_parameter_names(name); in ProgramShapeFromProto()
271 return std::move(program_shape); in ProgramShapeFromProto()
357 TF_ASSIGN_OR_RETURN(ProgramShape program_shape, in CreateModuleConfigFromProto()
360 HloModuleConfig module_config(program_shape); in CreateModuleConfigFromProto()
369 program_shape.parameters(i))); in CreateModuleConfigFromProto()
372 program_shape.result())); in CreateModuleConfigFromProto()
Dservice.cc265 const ProgramShape& program_shape, in CreateModuleConfig() argument
269 auto config = MakeUnique<HloModuleConfig>(program_shape); in CreateModuleConfig()
272 if (program_shape.parameters_size() != argument_shapes.size()) { in CreateModuleConfig()
274 program_shape.parameters_size(), in CreateModuleConfig()
281 program_shape.parameters(i))) { in CreateModuleConfig()
286 i, ShapeUtil::HumanString(program_shape.parameters(i)).c_str(), in CreateModuleConfig()
298 program_shape.result())); in CreateModuleConfig()
325 const ProgramShape& program_shape, in CreateModuleConfig() argument
333 return CreateModuleConfig(program_shape, argument_shapes, &execution_options, in CreateModuleConfig()
742 std::shared_ptr<const ProgramShape> program_shape, in ExecuteParallel()
[all …]
Dhlo_module_config.h41 explicit HloModuleConfig(const ProgramShape& program_shape);
50 void SetDefaultComputationLayout(const ProgramShape& program_shape);
Dcompile_only_service.cc96 std::shared_ptr<const ProgramShape> program_shape, in CompileAheadOfTime()
103 CreateModuleConfig(*program_shape, instance.argument_layouts, in CompileAheadOfTime()
Dhlo_computation.cc527 ProgramShape program_shape; in ComputeProgramShape() local
530 *program_shape.add_parameters() = param_instruction->shape(); in ComputeProgramShape()
531 *program_shape.add_parameter_names() = param_instruction->name(); in ComputeProgramShape()
533 *program_shape.mutable_result() = root_instruction_->shape(); in ComputeProgramShape()
535 LayoutUtil::ClearLayout(&program_shape); in ComputeProgramShape()
536 return program_shape; in ComputeProgramShape()
Dcomputation_layout.h38 explicit ComputationLayout(const ProgramShape& program_shape);
Dservice.h252 const ProgramShape& program_shape,
277 const ProgramShape& program_shape,
/external/tensorflow/tensorflow/compiler/xla/tools/
Ddumped_computation_to_operation_list.cc80 std::unique_ptr<ProgramShape> program_shape = in RealMain() local
84 layouts.reserve(program_shape->parameters_size()); in RealMain()
85 for (int i = 0; i < program_shape->parameters_size(); ++i) { in RealMain()
86 layouts.push_back(&program_shape->parameters(i)); in RealMain()
90 build_options.set_result_layout(program_shape->result()); in RealMain()
Ddumped_computation_to_text.cc55 std::unique_ptr<ProgramShape> program_shape = in RealMain() local
59 layouts.reserve(program_shape->parameters_size()); in RealMain()
60 for (int i = 0; i < program_shape->parameters_size(); ++i) { in RealMain()
61 layouts.push_back(&program_shape->parameters(i)); in RealMain()
66 build_options.set_result_layout(program_shape->result()); in RealMain()
/external/tensorflow/tensorflow/compiler/xla/tests/
Dhlo_test_base.cc76 ProgramShape program_shape; in GetProgramShapeWithLayout() local
79 *program_shape.add_parameters() = param->shape(); in GetProgramShapeWithLayout()
80 *program_shape.add_parameter_names() = param->name(); in GetProgramShapeWithLayout()
82 *program_shape.mutable_result() = entry->root_instruction()->shape(); in GetProgramShapeWithLayout()
83 return program_shape; in GetProgramShapeWithLayout()
128 const auto& program_shape = GetProgramShapeWithLayout(test_module); in MakeReferenceModule() local
132 if (!ProgramShapesEqual(program_shape, in MakeReferenceModule()
/external/tensorflow/tensorflow/compiler/xla/
Dlayout_util.cc135 /* static */ void LayoutUtil::SetToDefaultLayout(ProgramShape* program_shape) { in SetToDefaultLayout() argument
136 for (auto& parameter_shape : *program_shape->mutable_parameters()) { in SetToDefaultLayout()
139 LayoutUtil::SetToDefaultLayout(program_shape->mutable_result()); in SetToDefaultLayout()
237 /* static */ void LayoutUtil::ClearLayout(ProgramShape* program_shape) { in ClearLayout() argument
238 for (auto& parameter_shape : *program_shape->mutable_parameters()) { in ClearLayout()
241 LayoutUtil::ClearLayout(program_shape->mutable_result()); in ClearLayout()
322 /* static */ bool LayoutUtil::HasLayout(const ProgramShape& program_shape) { in HasLayout() argument
323 for (auto& parameter_shape : program_shape.parameters()) { in HasLayout()
328 return LayoutUtil::HasLayout(program_shape.result()); in HasLayout()
Dlayout_util.h61 static void SetToDefaultLayout(ProgramShape* program_shape);
76 static void ClearLayout(ProgramShape* program_shape);
130 static bool HasLayout(const ProgramShape& program_shape);
Dshape_util.cc174 ProgramShape program_shape; in MakeProgramShape() local
176 *program_shape.add_parameters() = shape; in MakeProgramShape()
178 *program_shape.mutable_result() = std::move(result); in MakeProgramShape()
179 return program_shape; in MakeProgramShape()
484 /* static */ string ShapeUtil::HumanString(const ProgramShape& program_shape) { in HumanString() argument
486 for (auto& shape : program_shape.parameters()) { in HumanString()
489 tensorflow::strings::StrCat(i < program_shape.parameter_names_size() in HumanString()
490 ? program_shape.parameter_names(i) in HumanString()
496 HumanString(program_shape.result())); in HumanString()
/external/tensorflow/tensorflow/compiler/aot/
Dcodegen.cc103 const xla::ProgramShape& ps = compile_result.program_shape; in ComputeArgSizes()
299 const xla::ProgramShape& ps = compile_result.program_shape; in GenerateHeader()
539 std::unique_ptr<xla::ProgramShape> program_shape; in GenerateMetadata() local
542 program_shape = in GenerateMetadata()
543 tensorflow::MakeUnique<xla::ProgramShape>(compile_result.program_shape); in GenerateMetadata()
547 program_shape->clear_parameter_names(); in GenerateMetadata()
557 "xla::ProgramShape", program_shape.get())); in GenerateMetadata()
Dcompile.cc57 compile_result->program_shape = *pshape_or.ValueOrDie(); in CompileXla()
58 xla::ProgramShape* pshape = &compile_result->program_shape; in CompileXla()
Dcompile.h36 xla::ProgramShape program_shape; // Static shape of args and results. member
/external/tensorflow/tensorflow/compiler/xla/client/lib/
Dtesting.cc101 auto program_shape = in MakeFakeArgumentsOrDie() local
107 for (const Shape& parameter : program_shape->parameters()) { in MakeFakeArgumentsOrDie()
/external/tensorflow/tensorflow/compiler/xla/python/
Dlocal_computation_builder.cc282 TF_ASSIGN_OR_RETURN(ProgramShape program_shape, in GetReturnValueShape()
284 return std::move(*program_shape.mutable_result()); in GetReturnValueShape()
313 TF_ASSIGN_OR_RETURN(ProgramShape program_shape, builder_.GetProgramShape()); in GetReturnValueShape()
314 return program_shape.result(); in GetReturnValueShape()

12