Home
last modified time | relevance | path

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

123

/external/tensorflow/tensorflow/compiler/xla/tests/
Dreplay_test.cc52 std::unique_ptr<ProgramShape> original_shape = in TEST_F()
54 std::unique_ptr<ProgramShape> replayed_shape = in TEST_F()
84 std::unique_ptr<ProgramShape> original_shape = in XLA_TEST_F()
86 std::unique_ptr<ProgramShape> replayed_shape = in XLA_TEST_F()
128 std::unique_ptr<ProgramShape> original_shape = in TEST_F()
130 std::unique_ptr<ProgramShape> replayed_shape = in TEST_F()
Daxpy_simple_test.cc65 TF_ASSERT_OK_AND_ASSIGN(ProgramShape shape, builder.GetProgramShape()); in TEST_F()
/external/tensorflow/tensorflow/compiler/xla/service/
Dshape_inference.h81 absl::Span<const Shape* const> arg_shapes, const ProgramShape& to_apply,
193 const ProgramShape& to_apply);
199 const ProgramShape& to_apply_shape);
206 const ProgramShape& to_apply_shape);
215 const Shape& operand_shape, const ProgramShape& select_shape,
217 const Shape& init_value_shape, const ProgramShape& scatter_shape);
257 static StatusOr<Shape> InferWhileShape(const ProgramShape& condition,
258 const ProgramShape& body,
264 absl::Span<const ProgramShape> branch_computations,
331 absl::Span<const Shape* const> arg_shapes, const ProgramShape& to_apply);
[all …]
Dhlo_module_config.cc31 HloModuleConfig::HloModuleConfig(const ProgramShape& program_shape, in HloModuleConfig()
40 const ProgramShape& program_shape) { in SetDefaultComputationLayout()
45 const ProgramShape& program_shape) { in SetComputationLayoutIfExists()
Dcomputation_layout.cc27 ComputationLayout::ComputationLayout(const ProgramShape& program_shape, in ComputationLayout()
74 ProgramShape ComputationLayout::ComputeProgramShape() const { in ComputeProgramShape()
75 ProgramShape program_shape; in ComputeProgramShape()
Dcomputation_layout.h42 explicit ComputationLayout(const ProgramShape& program_shape,
90 ProgramShape ComputeProgramShape() const;
Dhlo_module_config.h78 explicit HloModuleConfig(const ProgramShape& program_shape,
91 void SetDefaultComputationLayout(const ProgramShape& program_shape);
95 void SetComputationLayoutIfExists(const ProgramShape& program_shape);
Dshape_inference_test.cc68 ProgramShape to_apply = ShapeUtil::MakeProgramShape({f32_, f32_}, f32_); in ExpectInferredReduceShape()
103 ProgramShape select_program_shape_;
104 ProgramShape scatter_program_shape_;
331 ProgramShape to_apply = ShapeUtil::MakeProgramShape( in TEST_F()
361 ProgramShape select_program_shape_fail = in TEST_F()
372 ProgramShape select_program_shape_fail = ShapeUtil::MakeProgramShape( in TEST_F()
383 ProgramShape select_program_shape_fail = ShapeUtil::MakeProgramShape( in TEST_F()
394 ProgramShape select_program_shape_fail = ShapeUtil::MakeProgramShape( in TEST_F()
968 ProgramShape to_apply = ShapeUtil::MakeProgramShape({f32_}, s32_); in TEST_F()
1039 ProgramShape to_apply = ShapeUtil::MakeProgramShape({f32_}, f32_); in TEST_F()
[all …]
Dhlo_module_util.h37 const ProgramShape& program_shape,
Dservice.h188 const ProgramShape& program_shape,
197 const ProgramShape& program_shape,
/external/tensorflow/tensorflow/compiler/xla/
Dshape.cc222 ProgramShape::ProgramShape(const ProgramShapeProto& program_shape_proto) { in ProgramShape() function in xla::ProgramShape
232 ProgramShapeProto ProgramShape::ToProto() const { in ToProto()
244 std::string ProgramShape::ToString() const { in ToString()
255 std::ostream& operator<<(std::ostream& out, const ProgramShape& program_shape) { in operator <<()
Dlayout_util.h74 static void SetToDefaultLayout(ProgramShape* program_shape);
92 static void ClearLayout(ProgramShape* program_shape);
152 static bool HasLayout(const ProgramShape& program_shape);
Dshape.h313 class ProgramShape {
315 ProgramShape() = default;
318 explicit ProgramShape(const ProgramShapeProto& program_shape_proto);
388 std::ostream& operator<<(std::ostream& out, const ProgramShape& program_shape);
Dshape_test.cc146 ProgramShape program_shape; in TEST_F()
163 ProgramShape program_shape_copy(program_shape.ToProto()); in TEST_F()
183 ProgramShape prog = ShapeUtil::MakeProgramShape( in TEST_F()
/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.cc28 const ProgramShape& program_shape, in CreateModuleConfig()
/external/tensorflow/tensorflow/compiler/tf2xla/
Dxla_jit_compiled_cpu_function_test.cc219 ASSERT_TRUE(function.ProgramShape() != nullptr); in TEST()
220 const xla::ProgramShape program_shape(*function.ProgramShape()); in TEST()
270 ASSERT_TRUE(function.ProgramShape() != nullptr); in TEST()
271 const xla::ProgramShape program_shape(*function.ProgramShape()); 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()
64 TF_ASSERT_OK_AND_ASSIGN(ProgramShape program_shape, in TEST_F()
/external/tensorflow/tensorflow/compiler/xla/python/
Dxla_compiler.cc363 py::class_<ProgramShape>(m, "ProgramShape") in BuildXlaCompilerSubmodule()
365 [](absl::Span<const Shape> params, Shape result) -> ProgramShape { in BuildXlaCompilerSubmodule()
366 ProgramShape program_shape; in BuildXlaCompilerSubmodule()
374 static_cast<const std::vector<Shape>& (ProgramShape::*)() const>( in BuildXlaCompilerSubmodule()
375 &ProgramShape::parameters)) in BuildXlaCompilerSubmodule()
376 .def("result_shape", &ProgramShape::result) in BuildXlaCompilerSubmodule()
377 .def("__repr__", &ProgramShape::ToString); in BuildXlaCompilerSubmodule()
535 std::optional<XlaOp> root) -> StatusOr<ProgramShape> { in BuildXlaCompilerSubmodule()
Dxla_client.py222 ProgramShape = _xla.ProgramShape variable
223 ProgramShape.__doc__ = """
/external/tensorflow/tensorflow/core/tpu/kernels/
Dtpu_compile_op_support.h97 const xla::ProgramShape& program_shape,
109 const xla::ProgramShape& program_shape,
Dtpu_compile_op_support.cc39 using ::xla::ProgramShape;
59 const ProgramShape& program_shape, absl::Span<const Shape> argument_shapes, in CreateModuleConfig()
137 const xla::ProgramShape& program_shape, in CreateModuleConfig()
325 xla::ProgramShape( in CreateHloModules()
/external/tensorflow/tensorflow/compiler/xla/tools/
Dshow_signature.cc54 std::unique_ptr<ProgramShape> shape = in RealMain()

123