Searched refs:device_shape (Results 1 – 2 of 2) sorted by relevance
342 xla::Shape* device_shape) { in ExpandTreeOfTuples() argument346 *device_shape = elements.shape(); in ExpandTreeOfTuples()366 *xla::ShapeUtil::GetMutableSubshape(device_shape, index) = in ExpandTreeOfTuples()389 xla::Shape device_shape; in MakeTuple() local391 &host_shape, &device_shape)); in MakeTuple()397 host_shape, device_shape); in MakeTuple()412 xla::ShapeUtil::GetSubshape(device_shape, index); in MakeTuple()
214 xla::Shape* device_shape);