Searched refs:return_shape (Results 1 – 4 of 4) sorted by relevance
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 …]
944 Shape return_shape; in TransferToClient() local946 return_shape = Shape(arg->shape_with_layout()); in TransferToClient()947 if (!LayoutUtil::HasLayout(return_shape)) { in TransferToClient()951 return_shape = Shape(shaped_buffer->on_host_shape()); in TransferToClient()962 if (LayoutUtil::LayoutsInShapesEqual(return_shape, result_literal.shape())) { in TransferToClient()966 result_literal.Relayout(return_shape).ToProto(); in TransferToClient()
1088 xla::Shape return_shape = xla::ShapeUtil::MakeTupleShape( in TEST() local1115 *shapes->mutable_result() = return_shape.ToProto(); in TEST()
3044 const Shape& return_shape = callee.root_instruction()->shape(); in EmitThreadLocalCall() local3049 CHECK(ShapeUtil::IsScalar(return_shape)); in EmitThreadLocalCall()3051 PrimitiveType return_type = return_shape.element_type(); in EmitThreadLocalCall()