Searched refs:GetProgramShape (Results 1 – 15 of 15) sorted by relevance
26 StatusOr<ProgramShape> XlaComputation::GetProgramShape() const { in GetProgramShape() function in xla::XlaComputation
46 StatusOr<ProgramShape> GetProgramShape() const;
154 StatusOr<ProgramShape> XlaBuilder::GetProgramShape(int64 root_id) const { in GetProgramShape() function in xla::XlaBuilder185 StatusOr<ProgramShape> XlaBuilder::GetProgramShape() const { in GetProgramShape() function in xla::XlaBuilder187 return GetProgramShape(instructions_.back().id()); in GetProgramShape()190 StatusOr<ProgramShape> XlaBuilder::GetProgramShape(XlaOp root) const { in GetProgramShape() function in xla::XlaBuilder194 return GetProgramShape(root.handle()); in GetProgramShape()351 TF_ASSIGN_OR_RETURN(ProgramShape program_shape, GetProgramShape(root_id)); in Build()620 computation.GetProgramShape()); in Call()1752 computation.GetProgramShape()); in Map()1825 TF_ASSIGN_OR_RETURN(const auto& body_program_shape, body.GetProgramShape()); in While()1827 condition.GetProgramShape()); in While()[all …]
446 TF_ASSIGN_OR_RETURN(const auto& result, computation.GetProgramShape()); in GetComputationShape()
251 StatusOr<ProgramShape> GetProgramShape() const;255 StatusOr<ProgramShape> GetProgramShape(XlaOp root) const;626 StatusOr<ProgramShape> GetProgramShape(int64 root_id) const;
65 TF_ASSERT_OK_AND_ASSIGN(ProgramShape shape, builder.GetProgramShape()); in TEST_F()
400 StatusOr<ProgramShape> Computation::GetProgramShape() const { in GetProgramShape() function in xla::swig::Computation401 return computation_.GetProgramShape(); in GetProgramShape()405 TF_ASSIGN_OR_RETURN(ProgramShape shape, computation_.GetProgramShape()); in GetReturnValueShape()442 TF_ASSIGN_OR_RETURN(ProgramShape program_shape, builder_.GetProgramShape()); in GetReturnValueShape()
178 StatusOr<ProgramShape> GetProgramShape() const;
278 %unignore xla::swig::Computation::GetProgramShape;
794 def GetProgramShape(self): member in Computation795 return _wrap_program_shape(self._c_computation.GetProgramShape())
246 xla_computation.GetProgramShape()); in ExpandInstruction()
494 xla_computation.GetProgramShape()); in ExpandInstruction()
300 TF_ASSIGN_OR_RETURN(ProgramShape program_shape, comparator.GetProgramShape()); in MakeSortHlo()
324 computation->GetProgramShape()); in BuildComputation()
813 auto result_shape = xla_computation.GetProgramShape().ValueOrDie().result(); in TEST()