Home
last modified time | relevance | path

Searched refs:GetProgramShape (Results 1 – 25 of 29) sorted by relevance

12

/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/cpu/tests/
H A Dcpu_topk_test.cc38 xla_computation.GetProgramShape()); in TEST_F()
65 xla_computation.GetProgramShape()); in TEST_F()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/client/lib/
H A Ddynamic_shaped_ops.cc128 auto true_shape = true_computation.GetProgramShape().value().result(); in DynamicConditional()
130 auto false_shape = false_computation.GetProgramShape().value().result(); in DynamicConditional()
175 branch_computations[i]->GetProgramShape()); in DynamicConditional()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/client/
H A Dxla_computation.cc26 StatusOr<ProgramShape> XlaComputation::GetProgramShape() const { in GetProgramShape() function in xla::XlaComputation
H A Dxla_builder.cc121 fused_computation.GetProgramShape()); in BuildFusion()
313 StatusOr<ProgramShape> XlaBuilder::GetProgramShape(int64_t root_id) const { in GetProgramShape() function in xla::XlaBuilder
344 StatusOr<ProgramShape> XlaBuilder::GetProgramShape() const { in GetProgramShape() function in xla::XlaBuilder
346 return GetProgramShape(instructions_.back().id()); in GetProgramShape()
349 StatusOr<ProgramShape> XlaBuilder::GetProgramShape(XlaOp root) const { in GetProgramShape() function in xla::XlaBuilder
353 return GetProgramShape(root.handle()); in GetProgramShape()
546 TF_ASSIGN_OR_RETURN(ProgramShape program_shape, GetProgramShape(root_id)); in Build()
898 computation.GetProgramShape()); in Call()
2252 computation.GetProgramShape()); in Map()
2372 TF_ASSIGN_OR_RETURN(const auto& body_program_shape, body.GetProgramShape()); in While()
[all …]
H A Dxla_computation.h46 StatusOr<ProgramShape> GetProgramShape() const;
H A Dclient.cc447 TF_ASSIGN_OR_RETURN(const auto& result, computation.GetProgramShape()); in GetComputationShape()
H A Dxla_builder.h340 StatusOr<ProgramShape> GetProgramShape() const;
344 StatusOr<ProgramShape> GetProgramShape(XlaOp root) const;
1005 StatusOr<ProgramShape> GetProgramShape(int64_t root_id) const;
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/tests/
H A Daxpy_simple_test.cc65 TF_ASSERT_OK_AND_ASSIGN(ProgramShape shape, builder.GetProgramShape()); in TEST_F()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/python/
H A Dxla_compiler.cc403 .def("program_shape", &XlaComputation::GetProgramShape) in BuildXlaCompilerSubmodule()
536 return root ? builder.GetProgramShape(*root) in BuildXlaCompilerSubmodule()
537 : builder.GetProgramShape(); in BuildXlaCompilerSubmodule()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/
H A Drng_bit_generator_expander.cc90 xla_computation.GetProgramShape()); in GetGeneratorComputation()
H A Dbitcast_dtypes_expander.cc108 xla_computation.GetProgramShape()); in ExpandInstruction()
H A Drng_expander.cc116 xla_computation.GetProgramShape()); in GetComputationForRng()
H A Dcholesky_expander.cc252 xla_computation.GetProgramShape()); in ExpandInstruction()
H A Deigh_expander.cc587 xla_computation.GetProgramShape()); in ExpandInstruction()
H A Dqr_expander.cc555 xla_computation.GetProgramShape()); in ExpandInstruction()
H A Dtriangular_solve_expander.cc597 xla_computation.GetProgramShape()); in ExpandInstruction()
H A Dhlo_creation_utils.cc545 TF_ASSIGN_OR_RETURN(ProgramShape program_shape, comparator.GetProgramShape()); in MakeSortHlo()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/jit/
H A Dget_compiler_ir.cc169 result.computation->GetProgramShape()); in GetCompilerIr()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/pjrt/
H A Dutils.cc156 computation.GetProgramShape()); in DetermineArgumentLayoutsFromCompileOptions()
H A Dtpu_client.cc182 computation.GetProgramShape()); in DeserializeExecutable()
H A Dtfrt_cpu_pjrt_client.cc252 computation.GetProgramShape()); in JitCompile()
386 computation.GetProgramShape()); in Compile()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/
H A Dtf_xla_mlir_translate.cc85 compilation_result.computation->GetProgramShape().ValueOrDie()); in PrintHloModuleText()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/gpu/
H A Dcudnn_vectorize_convolutions.cc79 TF_ASSIGN_OR_RETURN(ProgramShape program_shape, comp.GetProgramShape()); in BuilderToHloComputation()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/spmd/
H A Dcustom_call_handler.cc191 TF_ASSIGN_OR_RETURN(ProgramShape program_shape, comparator.GetProgramShape()); in HandleCustomCallTopK()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/tf2xla/
H A Dxla_compiler_test.cc429 EXPECT_EQ(result.computation->GetProgramShape().value().result(), in TEST_F()
1191 EXPECT_EQ(result.computation->GetProgramShape().value().result(), in TEST_F()
1236 EXPECT_EQ(result.computation->GetProgramShape().value().result(), in TEST_F()

12