Searched refs:TakeSubTree (Results 1 – 5 of 5) sorted by relevance
134 ScopedShapedBuffer output = sb.TakeSubTree(subtree_index); in TEST()191 (void)shaped_buffer.TakeSubTree(/*index=*/{fan_out / 2}).release(); in BM_TakeSubTree()
201 ScopedShapedBuffer TakeSubTree(ShapeIndexView index);
175 ScopedShapedBuffer ScopedShapedBuffer::TakeSubTree(ShapeIndexView index) { in TakeSubTree() function in xla::ScopedShapedBuffer
336 xla_tensor->set_shaped_buffer(output->TakeSubTree({output_num})); in PopulateXlaTensor()
1835 ScopedShapedBuffer tuple_buffer = result_buffer.TakeSubTree({i}); in MakeOutputBuffers()