Home
last modified time | relevance | path

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

123

/external/tensorflow/tensorflow/compiler/xla/
Dshape_test.cc146 ProgramShape program_shape; in TEST_F() local
147 *program_shape.add_parameters() = ShapeUtil::MakeShape(F32, {1, 2, 3}); in TEST_F()
148 *program_shape.add_parameters() = ShapeUtil::MakeTokenShape(); in TEST_F()
149 *program_shape.add_parameters() = ShapeUtil::MakeShape(S64, {}); in TEST_F()
150 *program_shape.add_parameters() = ShapeUtil::MakeTupleShape( in TEST_F()
155 *program_shape.mutable_result() = ShapeUtil::MakeShape(F32, {7}); in TEST_F()
157 program_shape.add_parameter_names("foo"); in TEST_F()
158 program_shape.add_parameter_names("bar"); in TEST_F()
159 program_shape.add_parameter_names("baz"); in TEST_F()
160 program_shape.add_parameter_names("qux qux"); in TEST_F()
[all …]
Dlayout_util.cc168 /* static */ void LayoutUtil::SetToDefaultLayout(ProgramShape* program_shape) { in SetToDefaultLayout() argument
169 for (auto& parameter_shape : *program_shape->mutable_parameters()) { in SetToDefaultLayout()
172 LayoutUtil::SetToDefaultLayout(program_shape->mutable_result()); in SetToDefaultLayout()
268 /* static */ void LayoutUtil::ClearLayout(ProgramShape* program_shape) { in ClearLayout() argument
269 for (auto& parameter_shape : *program_shape->mutable_parameters()) { in ClearLayout()
272 LayoutUtil::ClearLayout(program_shape->mutable_result()); in ClearLayout()
307 /* static */ bool LayoutUtil::HasLayout(const ProgramShape& program_shape) { in HasLayout() argument
308 for (auto& parameter_shape : program_shape.parameters()) { in HasLayout()
313 return LayoutUtil::HasLayout(program_shape.result()); in HasLayout()
Dlayout_util.h74 static void SetToDefaultLayout(ProgramShape* program_shape);
92 static void ClearLayout(ProgramShape* program_shape);
117 static bool HasLayout(const ProgramShape& program_shape);
/external/tensorflow/tensorflow/compiler/xla/service/
Dcomputation_layout.cc27 ComputationLayout::ComputationLayout(const ProgramShape& program_shape, in ComputationLayout() argument
29 : result_layout_(program_shape.result()) { in ComputationLayout()
30 for (auto& shape : program_shape.parameters()) { in ComputationLayout()
74 ProgramShape program_shape; in ComputeProgramShape() local
76 *program_shape.add_parameters() = parameter_layouts_[i].shape(); in ComputeProgramShape()
77 *program_shape.add_parameter_names() = absl::StrCat("p", i); in ComputeProgramShape()
79 *program_shape.mutable_result() = result_layout_.shape(); in ComputeProgramShape()
80 return program_shape; in ComputeProgramShape()
Dhlo_module_config.cc31 HloModuleConfig::HloModuleConfig(const ProgramShape& program_shape, in HloModuleConfig() argument
34 ComputationLayout(program_shape, ignore_layouts)) {} in HloModuleConfig()
40 const ProgramShape& program_shape) { in SetDefaultComputationLayout() argument
41 entry_computation_layout_ = ComputationLayout(program_shape); in SetDefaultComputationLayout()
45 const ProgramShape& program_shape) { in SetComputationLayoutIfExists() argument
46 entry_computation_layout_ = ComputationLayout(program_shape, in SetComputationLayoutIfExists()
Dhlo_module_util.cc45 const ProgramShape& program_shape, in CreateModuleConfig() argument
49 auto config = absl::make_unique<HloModuleConfig>(program_shape); in CreateModuleConfig()
53 if (program_shape.parameters_size() != argument_shapes_size) { in CreateModuleConfig()
55 program_shape.parameters_size(), in CreateModuleConfig()
62 program_shape.parameters(i))) { in CreateModuleConfig()
66 i, ShapeUtil::HumanString(program_shape.parameters(i)), in CreateModuleConfig()
78 ValidateResultShape(shape_with_output_layout, program_shape.result())); in CreateModuleConfig()
Dlocal_service.cc107 ProgramShape program_shape(proto.host_program_shape()); in CompileExecutables() local
110 if (argument_layouts.size() != program_shape.parameters_size()) { in CompileExecutables()
113 program_shape.parameters_size(), argument_layouts.size()); in CompileExecutables()
120 if (!ShapeUtil::Compatible(argument_shape, program_shape.parameters(i))) { in CompileExecutables()
137 ShapeUtil::HumanString(program_shape.parameters(i)), in CompileExecutables()
143 program_shape.result())); in CompileExecutables()
147 CreateExecutionOptions(build_options, &program_shape); in CompileExecutables()
151 CreateModuleConfig(program_shape, argument_layouts, &execution_options)); in CompileExecutables()
Dhlo_module_config.h74 explicit HloModuleConfig(const ProgramShape& program_shape,
87 void SetDefaultComputationLayout(const ProgramShape& program_shape);
91 void SetComputationLayoutIfExists(const ProgramShape& program_shape);
Dhlo_proto_util.cc64 const auto& program_shape = hlo_proto.hlo_module().host_program_shape(); in EntryComputationParameterShapes() local
65 for (const ShapeProto& shape : program_shape.parameters()) { in EntryComputationParameterShapes()
Drng_bit_generator_expander.cc89 TF_ASSIGN_OR_RETURN(ProgramShape program_shape, in GetGeneratorComputation()
91 HloModuleConfig config(program_shape); in GetGeneratorComputation()
/external/tensorflow/tensorflow/compiler/xla/pjrt/
Dutils.cc75 const XlaComputation& computation, const ProgramShape& program_shape) { in GetShardedProgramShapes() argument
77 arg_shapes.resize(program_shape.parameters_size()); in GetShardedProgramShapes()
85 if (instr.parameter_number() >= program_shape.parameters_size()) { in GetShardedProgramShapes()
88 instr.parameter_number(), program_shape.parameters_size()); in GetShardedProgramShapes()
153 TF_ASSIGN_OR_RETURN(ProgramShape program_shape, in DetermineArgumentLayoutsFromCompileOptions()
156 argument_layouts.emplace(program_shape.parameters()); in DetermineArgumentLayoutsFromCompileOptions()
160 } else if (argument_layouts->size() != program_shape.parameters_size()) { in DetermineArgumentLayoutsFromCompileOptions()
164 argument_layouts->size(), program_shape.parameters_size()); in DetermineArgumentLayoutsFromCompileOptions()
188 GetShardedProgramShapes(computation, program_shape)); in DetermineArgumentLayoutsFromCompileOptions()
201 result_layout = program_shape.result(); in DetermineArgumentLayoutsFromCompileOptions()
/external/tensorflow/tensorflow/compiler/xla/tools/
Ddumped_computation_to_operation_list.cc77 std::unique_ptr<ProgramShape> program_shape = in RealMain() local
81 layouts.reserve(program_shape->parameters_size()); in RealMain()
82 for (int i = 0; i < program_shape->parameters_size(); ++i) { in RealMain()
83 layouts.push_back(&program_shape->parameters(i)); in RealMain()
87 build_options.set_result_layout(program_shape->result()); in RealMain()
Ddumped_computation_to_text.cc53 std::unique_ptr<ProgramShape> program_shape = in RealMain() local
57 layouts.reserve(program_shape->parameters_size()); in RealMain()
58 for (int i = 0; i < program_shape->parameters_size(); ++i) { in RealMain()
59 layouts.push_back(&program_shape->parameters(i)); in RealMain()
64 build_options.set_result_layout(program_shape->result()); in RealMain()
/external/tensorflow/tensorflow/compiler/tf2xla/
Dxla_jit_compiled_cpu_function.cc97 TF_ASSIGN_OR_RETURN(std::unique_ptr<xla::ProgramShape> program_shape, in Compile()
99 if (program_shape->result().element_type() != xla::TUPLE) { in Compile()
108 program_shape->clear_parameter_names(); in Compile()
112 arg_shapes.reserve(program_shape->parameters_size()); in Compile()
113 for (int i = 0; i < program_shape->parameters_size(); ++i) { in Compile()
114 arg_shapes.push_back(&program_shape->parameters(i)); in Compile()
146 absl::make_unique<xla::ProgramShapeProto>(program_shape->ToProto()); in Compile()
Dxla_jit_compiled_cpu_function_test.cc220 const xla::ProgramShape program_shape(*function.ProgramShape()); in TEST() local
221 ASSERT_EQ(program_shape.parameters_size(), 2); in TEST()
222 EXPECT_TRUE(ShapeUtil::Compatible(program_shape.parameters(0), s32)); in TEST()
223 EXPECT_TRUE(ShapeUtil::Compatible(program_shape.parameters(1), s32)); in TEST()
225 const xla::Shape& result = program_shape.result(); in TEST()
271 const xla::ProgramShape program_shape(*function.ProgramShape()); in TEST() local
272 ASSERT_EQ(program_shape.parameters_size(), 2); in TEST()
273 EXPECT_TRUE(ShapeUtil::Compatible(program_shape.parameters(0), s32)); in TEST()
274 EXPECT_TRUE(ShapeUtil::Compatible(program_shape.parameters(1), s32_1)); in TEST()
276 const xla::Shape& result = program_shape.result(); in TEST()
/external/tensorflow/tensorflow/compiler/xla/service/cpu/tests/
Dcpu_topk_test.cc37 TF_ASSERT_OK_AND_ASSIGN(ProgramShape program_shape, in TEST_F()
39 HloModuleConfig config(program_shape); in TEST_F()
64 TF_ASSERT_OK_AND_ASSIGN(ProgramShape program_shape, in TEST_F()
66 HloModuleConfig config(program_shape); in TEST_F()
/external/tensorflow/tensorflow/compiler/xrt/kernels/
Dxrt_compile_ops.cc153 config.program_shape().parameters_size()); in Compile()
155 config.program_shape().parameters_size()); in Compile()
156 for (int i = 0; i < config.program_shape().parameters_size(); ++i) { in Compile()
157 argument_layouts[i] = xla::Shape(config.program_shape().parameters(i)); in Compile()
163 build_options.set_result_layout(xla::Shape(config.program_shape().result())); in Compile()
229 xla::ProgramShapeProto program_shape = executable->executable() in Compute() local
236 program_shape_output.vec<tstring>()(0) = program_shape.SerializeAsString(); in Compute()
/external/tensorflow/tensorflow/core/tpu/kernels/
Dtpu_compile_op_support.cc57 const ProgramShape& program_shape, absl::Span<const Shape> argument_shapes, in CreateModuleConfig() argument
64 auto config = absl::make_unique<HloModuleConfig>(program_shape); in CreateModuleConfig()
67 if (program_shape.parameters_size() != argument_shapes.size()) { in CreateModuleConfig()
69 program_shape.parameters_size(), in CreateModuleConfig()
76 program_shape.parameters(i))) { in CreateModuleConfig()
80 i, ShapeUtil::HumanString(program_shape.parameters(i)), in CreateModuleConfig()
90 ValidateResultShape(result_layout.value(), program_shape.result())); in CreateModuleConfig()
135 const xla::ProgramShape& program_shape, in CreateModuleConfig() argument
140 return CreateModuleConfig(program_shape, argument_shapes, result_layout, in CreateModuleConfig()
Dtpu_compile_op_support.h96 const xla::ProgramShape& program_shape,
108 const xla::ProgramShape& program_shape,
/external/tensorflow/tensorflow/compiler/xla/client/
Dcompile_only_client.cc27 const ProgramShape& program_shape, in CreateModuleConfig() argument
30 return compiler_service_->CreateModuleConfig(program_shape, argument_shapes, in CreateModuleConfig()
Dexecutable_build_options.cc111 const ProgramShape* program_shape) { in CreateExecutionOptions() argument
120 Shape result_shape(program_shape->result()); in CreateExecutionOptions()
/external/tensorflow/tensorflow/compiler/xla/client/lib/
Dtesting.cc102 auto program_shape = computation.proto().host_program_shape(); in MakeFakeArgumentsOrDie() local
105 for (const ShapeProto& shape : program_shape.parameters()) { in MakeFakeArgumentsOrDie()
/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()
250 const auto& program_shape = GetProgramShapeWithLayout(test_module); in MakeReferenceModule() local
254 if (!ProgramShapesEqual(program_shape, in MakeReferenceModule()
/external/tensorflow/tensorflow/compiler/aot/
Dcodegen.cc417 const xla::ProgramShapeProto& ps = compile_result.program_shape; in GenerateHeader()
747 std::unique_ptr<xla::ProgramShapeProto> program_shape; in GenerateMetadata() local
750 program_shape = in GenerateMetadata()
751 absl::make_unique<xla::ProgramShapeProto>(compile_result.program_shape); in GenerateMetadata()
756 program_shape->clear_parameter_names(); in GenerateMetadata()
764 "::xla::ProgramShapeProto", program_shape.get()}; in GenerateMetadata()
/external/tensorflow/tensorflow/core/tpu/
Dtpu_embedding_optimization_parameters_utils.cc154 const xla::ProgramShapeProto& program_shape = in GetBaseAuxiliaryParameterCount() local
157 const int num_inputs = program_shape.parameters_size(); in GetBaseAuxiliaryParameterCount()
158 const int num_outputs = program_shape.result().tuple_shapes_size(); in GetBaseAuxiliaryParameterCount()

123