Searched defs:shape_representation_fn (Results 1 – 6 of 6) sorted by relevance
24 const ShapeRepresentationFn& shape_representation_fn, in RewriteLayoutWithShardedShape()69 const ShapeRepresentationFn& shape_representation_fn, in ReshapeWithCorrectRepresentationAndSharding()
48 ShapeRepresentationFn shape_representation_fn; member
55 XlaHelpers::ShapeRepresentationFn shape_representation_fn; member
150 HloModule* module, DeviceShapeRepresentationFn shape_representation_fn, in UpdateEntryComputationLayout()
383 int device_ordinal, DeviceShapeRepresentationFn shape_representation_fn) { in AllocateScopedShapedBuffer()
300 XlaHelpers::ShapeRepresentationFn shape_representation_fn = in CompileMlirToXlaHloViaBuilder() local