Searched refs:use_spmd_attr (Results 1 – 3 of 3) sorted by relevance
76 if (auto use_spmd_attr = cluster_func->getAttrOfType<BoolAttr>( in PartitionResourceReadsWrites() local78 use_spmd = use_spmd_attr.getValue(); in PartitionResourceReadsWrites()
316 if (auto use_spmd_attr = cluster_func->getAttrOfType<BoolAttr>( in IdentifyXlaShardingForTPUComputation() local318 use_spmd = use_spmd_attr.getValue(); in IdentifyXlaShardingForTPUComputation()
317 auto use_spmd_attr = op->getAttrOfType<BoolAttr>(kUseXlaSpmdAttr); in SetMetadataProtoFromClusterFuncOp() local318 if (!use_spmd_attr) in SetMetadataProtoFromClusterFuncOp()320 metadata->set_use_spmd_for_xla_partitioning(use_spmd_attr.getValue()); in SetMetadataProtoFromClusterFuncOp()