Home
last modified time | relevance | path

Searched refs:enclosing_cluster (Results 1 – 3 of 3) sorted by relevance

/aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/
H A Dlayout_parsing.cc137 auto enclosing_cluster = op->getParentOfType<mlir::tf_device::ClusterOp>(); in ExtractDeviceMeshEnclosingCluster() local
138 if (!enclosing_cluster) in ExtractDeviceMeshEnclosingCluster()
141 TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshFromOp(enclosing_cluster)); in ExtractDeviceMeshEnclosingCluster()
H A Dspmd_expander_common.cc349 auto enclosing_cluster = in TraceUseToNextTFOp() local
352 enclosing_cluster.getResult(operand->getOperandNumber())); in TraceUseToNextTFOp()
/aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/expansions/
H A Dmeta_spmd_expander.cc1165 auto enclosing_cluster = op->getParentOfType<mlir::tf_device::ClusterOp>(); in ExpandOp() local
1166 if (!enclosing_cluster) in ExpandOp()
1207 if (use.getOwner() == enclosing_cluster.GetBody().getTerminator()) { in ExpandOp()
1219 enclosing_cluster->setAttr(kShapeOpInputLayoutIndices, in ExpandOp()
1221 enclosing_cluster->setAttr( in ExpandOp()