Searched defs:sub_shape (Results 1 – 9 of 9) sorted by relevance
59 const Shape& sub_shape = ShapeUtil::GetSubshape(literal_shape, leaf.first); in TransferLiteralToInfeed() local
392 const ShapeIndex& /*index*/) { in GetIndexTypeForKernel()440 const ShapeIndex& /*index*/) { in GetIndexTypeForKernelFromMlir()
99 [&](const Shape& sub_shape, const ShapeIndex& shape_index) { in MarkLiveAtAllIndices()
629 const Shape* sub_shape = &shape; in GetSubSharding() local
136 xla::Shape* sub_shape = in GetInfeedShapeWithLayout() local
219 const xla::ShapeIndex& index) -> Status { in UpdateDynamicInputs()
963 ForEachSubshape(shape, [&](const Shape& sub_shape, const ShapeIndex& index) { in GetLeafShapes()1592 const ShapeIndex& index) { in DynamicShapeIsCompatible()
1427 const Shape& sub_shape = in GetOrCreateView() local
1970 SmallVector<int64_t, 2> sub_shape; in fold() local