Home
last modified time | relevance | path

Searched refs:ProgramShape (Results 1 – 25 of 53) sorted by relevance

123

/external/tensorflow/tensorflow/compiler/xla/service/
Dshape_inference.h83 absl::Span<const Shape* const> arg_shapes, const ProgramShape& to_apply,
154 const ProgramShape& to_apply);
160 const ProgramShape& to_apply_shape);
169 const Shape& operand_shape, const ProgramShape& select_shape,
171 const Shape& init_value_shape, const ProgramShape& scatter_shape);
211 static StatusOr<Shape> InferWhileShape(const ProgramShape& condition,
212 const ProgramShape& body,
218 absl::Span<const ProgramShape> branch_computations,
276 absl::Span<const Shape* const> arg_shapes, const ProgramShape& to_apply);
297 const Shape& updates_shape, const ProgramShape& to_apply_shape,
Dcomputation_layout.cc27 ComputationLayout::ComputationLayout(const ProgramShape& program_shape, in ComputationLayout()
60 ProgramShape ComputationLayout::ComputeProgramShape() const { in ComputeProgramShape()
61 ProgramShape program_shape; in ComputeProgramShape()
Dhlo_module_config.cc31 HloModuleConfig::HloModuleConfig(const ProgramShape& program_shape, in HloModuleConfig()
40 const ProgramShape& program_shape) { in SetDefaultComputationLayout()
Dcomputation_layout.h42 explicit ComputationLayout(const ProgramShape& program_shape,
88 ProgramShape ComputeProgramShape() const;
Dhlo_module_config.h74 explicit HloModuleConfig(const ProgramShape& program_shape,
86 void SetDefaultComputationLayout(const ProgramShape& program_shape);
Dshape_inference_test.cc65 ProgramShape to_apply = ShapeUtil::MakeProgramShape({f32_, f32_}, f32_); in ExpectInferredReduceShape()
100 ProgramShape select_program_shape_;
101 ProgramShape scatter_program_shape_;
328 ProgramShape to_apply = ShapeUtil::MakeProgramShape( in TEST_F()
358 ProgramShape select_program_shape_fail = in TEST_F()
369 ProgramShape select_program_shape_fail = ShapeUtil::MakeProgramShape( in TEST_F()
380 ProgramShape select_program_shape_fail = ShapeUtil::MakeProgramShape( in TEST_F()
391 ProgramShape select_program_shape_fail = ShapeUtil::MakeProgramShape( in TEST_F()
757 ProgramShape to_apply = ShapeUtil::MakeProgramShape({f32_}, s32_); in TEST_F()
828 ProgramShape to_apply = ShapeUtil::MakeProgramShape({f32_}, f32_); in TEST_F()
[all …]
Dservice.h189 const ProgramShape& program_shape,
198 const ProgramShape& program_shape,
Dlocal_service.cc99 const ProgramShape* program_shape) { in CreateExecutionOptions()
129 ProgramShape program_shape(proto.host_program_shape()); in CompileExecutables()
Dhlo_module.cc307 ProgramShape expected_program_shape(proto.host_program_shape()); in CreateFromProto()
396 const ProgramShape& program_shape, const DebugOptions& debug_options, in CreateModuleConfigFromShape()
398 HloModuleConfig module_config(ProgramShape{program_shape}); in CreateModuleConfigFromShape()
443 ProgramShape program_shape(module.host_program_shape()); in CreateModuleConfigFromProto()
/external/tensorflow/tensorflow/compiler/xla/tests/
Dreplay_test.cc54 std::unique_ptr<ProgramShape> original_shape = in TEST_F()
56 std::unique_ptr<ProgramShape> replayed_shape = in TEST_F()
87 std::unique_ptr<ProgramShape> original_shape = in XLA_TEST_F()
89 std::unique_ptr<ProgramShape> replayed_shape = in XLA_TEST_F()
134 std::unique_ptr<ProgramShape> original_shape = in TEST_F()
136 std::unique_ptr<ProgramShape> replayed_shape = in TEST_F()
Dhlo_test_base.cc63 bool ProgramShapesEqual(const ProgramShape& lhs, const ProgramShape& rhs) { in ProgramShapesEqual()
75 ProgramShape GetProgramShapeWithLayout(const HloModule& module) { in GetProgramShapeWithLayout()
76 ProgramShape program_shape; in GetProgramShapeWithLayout()
Daxpy_simple_test.cc65 TF_ASSERT_OK_AND_ASSIGN(ProgramShape shape, builder.GetProgramShape()); in TEST_F()
/external/tensorflow/tensorflow/compiler/xla/
Dshape.cc189 ProgramShape::ProgramShape(const ProgramShapeProto& program_shape_proto) { in ProgramShape() function in xla::ProgramShape
199 ProgramShapeProto ProgramShape::ToProto() const { in ToProto()
211 string ProgramShape::ToString() const { in ToString()
222 std::ostream& operator<<(std::ostream& out, const ProgramShape& program_shape) { in operator <<()
Dlayout_util.h70 static void SetToDefaultLayout(ProgramShape* program_shape);
88 static void ClearLayout(ProgramShape* program_shape);
113 static bool HasLayout(const ProgramShape& program_shape);
Dshape.h255 class ProgramShape {
257 ProgramShape() = default;
260 explicit ProgramShape(const ProgramShapeProto& program_shape_proto);
328 std::ostream& operator<<(std::ostream& out, const ProgramShape& program_shape);
Dshape_test.cc145 ProgramShape program_shape; in TEST_F()
162 ProgramShape program_shape_copy(program_shape.ToProto()); in TEST_F()
182 ProgramShape prog = ShapeUtil::MakeProgramShape( in TEST_F()
Dlayout_util.cc162 /* static */ void LayoutUtil::SetToDefaultLayout(ProgramShape* program_shape) { in SetToDefaultLayout()
262 /* static */ void LayoutUtil::ClearLayout(ProgramShape* program_shape) { in ClearLayout()
301 /* static */ bool LayoutUtil::HasLayout(const ProgramShape& program_shape) { in HasLayout()
/external/tensorflow/tensorflow/compiler/xla/client/
Dxla_computation.cc26 StatusOr<ProgramShape> XlaComputation::GetProgramShape() const { in GetProgramShape()
28 return ProgramShape(proto_.host_program_shape()); in GetProgramShape()
Dxla_computation.h46 StatusOr<ProgramShape> GetProgramShape() const;
Dcompile_only_client.h62 const ProgramShape& program_shape,
Dcompile_only_client.cc26 const ProgramShape& program_shape, in CreateModuleConfig()
/external/tensorflow/tensorflow/compiler/tf2xla/
Dxla_jit_compiled_cpu_function_test.cc216 ASSERT_TRUE(function.ProgramShape() != nullptr); in TEST()
217 const xla::ProgramShape program_shape(*function.ProgramShape()); in TEST()
259 ASSERT_TRUE(function.ProgramShape() != nullptr); in TEST()
260 const xla::ProgramShape program_shape(*function.ProgramShape()); in TEST()
/external/tensorflow/tensorflow/compiler/xla/tools/
Dshow_signature.cc54 std::unique_ptr<ProgramShape> shape = in RealMain()
Ddumped_computation_to_operation_list.cc77 std::unique_ptr<ProgramShape> program_shape = in RealMain()
/external/tensorflow/tensorflow/compiler/xla/python/
Dxla.cc436 py::class_<ProgramShape>(m, "ProgramShape") in PYBIND11_MODULE()
438 [](absl::Span<const Shape> params, Shape result) -> ProgramShape { in PYBIND11_MODULE()
439 ProgramShape program_shape; in PYBIND11_MODULE()
447 static_cast<const std::vector<Shape>& (ProgramShape::*)() const>( in PYBIND11_MODULE()
448 &ProgramShape::parameters)) in PYBIND11_MODULE()
449 .def("result_shape", &ProgramShape::result) in PYBIND11_MODULE()
450 .def("__repr__", &ProgramShape::ToString); in PYBIND11_MODULE()
799 absl::optional<XlaOp> root) -> StatusOr<ProgramShape> { in PYBIND11_MODULE()

123