Searched refs:GetProgramShape (Results 1 – 17 of 17) sorted by relevance
68 compilation_result.computation->GetProgramShape().ValueOrDie()); in TEST()140 compilation_result.computation->GetProgramShape().ValueOrDie()); in TEST()178 compilation_result.computation->GetProgramShape().ValueOrDie()); in TEST()
26 StatusOr<ProgramShape> XlaComputation::GetProgramShape() const { in GetProgramShape() function in xla::XlaComputation
46 StatusOr<ProgramShape> GetProgramShape() const;
201 StatusOr<ProgramShape> XlaBuilder::GetProgramShape(int64 root_id) const { in GetProgramShape() function in xla::XlaBuilder232 StatusOr<ProgramShape> XlaBuilder::GetProgramShape() const { in GetProgramShape() function in xla::XlaBuilder234 return GetProgramShape(instructions_.back().id()); in GetProgramShape()237 StatusOr<ProgramShape> XlaBuilder::GetProgramShape(XlaOp root) const { in GetProgramShape() function in xla::XlaBuilder241 return GetProgramShape(root.handle()); in GetProgramShape()404 TF_ASSIGN_OR_RETURN(ProgramShape program_shape, GetProgramShape(root_id)); in Build()724 computation.GetProgramShape()); in Call()1690 computation.GetProgramShape()); in Map()1761 TF_ASSIGN_OR_RETURN(const auto& body_program_shape, body.GetProgramShape()); in While()1763 condition.GetProgramShape()); in While()[all …]
446 TF_ASSIGN_OR_RETURN(const auto& result, computation.GetProgramShape()); in GetComputationShape()
284 StatusOr<ProgramShape> GetProgramShape() const;288 StatusOr<ProgramShape> GetProgramShape(XlaOp root) const;666 StatusOr<ProgramShape> GetProgramShape(int64 root_id) const;
65 TF_ASSERT_OK_AND_ASSIGN(ProgramShape shape, builder.GetProgramShape()); in TEST_F()
114 xla_computation.GetProgramShape()); in GetComputationForRng()
240 xla_computation.GetProgramShape()); in ExpandInstruction()
511 xla_computation.GetProgramShape()); in ExpandInstruction()
369 TF_ASSIGN_OR_RETURN(ProgramShape program_shape, comparator.GetProgramShape()); in MakeSortHlo()
591 def GetProgramShape(self): member in Computation592 return self._c_computation.GetProgramShape()595 return self._c_computation.GetProgramShape().result_shape()
775 .def("GetProgramShape", &XlaComputation::GetProgramShape) in PYBIND11_MODULE()800 return root ? builder.GetProgramShape(*root) in PYBIND11_MODULE()801 : builder.GetProgramShape(); in PYBIND11_MODULE()
1054 computation.GetProgramShape()); in Compile()1089 computation.GetProgramShape()); in Compile()
776 << computation.GetProgramShape().ValueOrDie().DebugString(); in Compile()
398 computation->GetProgramShape()); in BuildComputation()
1109 auto result_shape = xla_computation.GetProgramShape().ValueOrDie().result(); in TEST()