Searched refs:return_shape (Results 1 – 6 of 6) sorted by relevance
120 const Shape& return_shape = root_instruction->shape(); in CodegenNestedComputation() local125 if (ShapeUtil::IsScalar(return_shape)) { in CodegenNestedComputation()130 CHECK(return_shape.IsTuple()); in CodegenNestedComputation()131 llvm::Type* tuple_type = llvm_ir::ShapeToIrType(return_shape, module_); in CodegenNestedComputation()135 for (int i = 0; i < return_shape.tuple_shapes_size(); i++) { in CodegenNestedComputation()136 const Shape& element_shape = return_shape.tuple_shapes(i); in CodegenNestedComputation()
778 const Shape& return_shape = function->root_instruction()->shape(); in HandleReduce() local781 llvm_ir::ShapeToIrType(return_shape, module_); in HandleReduce()783 llvm_ir::IrArray tuple_array(ret_argument, return_shape); in HandleReduce()
783 const Shape* return_shape = &shape; in GetSubshape() local785 CHECK(return_shape->IsTuple()) in GetSubshape()787 return_shape = &return_shape->tuple_shapes(i); in GetSubshape()789 return *return_shape; in GetSubshape()794 const Shape* return_shape = &shape; in TryGetSubshape() local796 if (!return_shape->IsTuple() || i < 0 || in TryGetSubshape()797 i >= return_shape->tuple_shapes_size()) { in TryGetSubshape()802 return_shape = &return_shape->tuple_shapes(i); in TryGetSubshape()804 return return_shape; in TryGetSubshape()809 Shape* return_shape = shape; in GetMutableSubshape() local[all …]
965 Shape return_shape; in TransferToClient() local967 return_shape = Shape(arg->shape_with_layout()); in TransferToClient()968 if (!LayoutUtil::HasLayout(return_shape)) { in TransferToClient()972 return_shape = Shape(shaped_buffer->on_host_shape()); in TransferToClient()983 if (LayoutUtil::LayoutsInShapesEqual(return_shape, result_literal.shape())) { in TransferToClient()987 result_literal.Relayout(return_shape).ToProto(); in TransferToClient()
118 const Shape& return_shape = computation->root_instruction()->shape(); in EmitThreadLocalFunctionEpilogue() local120 if (ShapeUtil::IsScalar(return_shape)) { in EmitThreadLocalFunctionEpilogue()126 CHECK(return_shape.IsTuple()); in EmitThreadLocalFunctionEpilogue()128 llvm::Type* tuple_type = llvm_ir::ShapeToIrType(return_shape, module_); in EmitThreadLocalFunctionEpilogue()132 for (int i = 0; i < return_shape.tuple_shapes_size(); i++) { in EmitThreadLocalFunctionEpilogue()133 const Shape& element_shape = return_shape.tuple_shapes(i); in EmitThreadLocalFunctionEpilogue()3363 const Shape& return_shape = callee.root_instruction()->shape(); in EmitThreadLocalCall() local3364 bool is_scalar_return = ShapeUtil::IsScalar(return_shape); in EmitThreadLocalCall()3366 return_shape.IsTuple() && in EmitThreadLocalCall()3367 absl::c_all_of(return_shape.tuple_shapes(), [&](const Shape& shape) { in EmitThreadLocalCall()[all …]
1385 xla::Shape return_shape = xla::ShapeUtil::MakeTupleShape( in TEST() local1412 *shapes->mutable_result() = return_shape.ToProto(); in TEST()