Home
last modified time | relevance | path

Searched defs:layout_op (Results 1 – 12 of 12) sorted by relevance

/external/tensorflow/tensorflow/dtensor/mlir/
Dundo_merge_const_across_mesh.cc48 mlir::TF::DTensorLayout layout_op = in runOnOperation() local
63 for (auto& layout_op : unique_layout_ops) { in runOnOperation() local
Dshape_utils.cc43 if (auto layout_op = in ExtractGlobalInputShape() local
80 if (auto layout_op = mlir::dyn_cast<mlir::TF::DTensorLayout>(user)) { in ExtractGlobalOutputShape() local
250 auto layout_op = mlir::cast<mlir::TF::DTensorLayout>(value.getDefiningOp()); in GetGlobalShapeOfValueFromDTensorLayout() local
254 auto layout_op = in GetGlobalShapeOfValueFromDTensorLayout() local
Dpropagate_default_layout.cc51 auto layout_op = builder->create<mlir::TF::DTensorLayout>( in CreateDTensorLayoutOp() local
193 auto layout_op = builder.create<mlir::TF::DTensorLayout>( in runOnOperation() local
Dop_to_device_cluster.cc56 if (auto layout_op = llvm::dyn_cast<mlir::TF::DTensorLayout>(op)) { in WrapDeviceCluster() local
Dsparse_expansion.cc80 if (auto layout_op = mlir::dyn_cast<mlir::TF::DTensorLayout>( in RemoveUnusedSparseToDenseOps() local
Dspmd_expansion.cc329 if (auto layout_op = llvm::dyn_cast<mlir::TF::DTensorLayout>(op)) in ConductSPMDExpansion() local
362 for (auto layout_op : layout_ops) RemoveDTensorLayoutOp(layout_op); in RemoveDTensorLayoutOps() local
Dlayout_parsing.cc167 if (auto layout_op = llvm::dyn_cast<mlir::TF::DTensorLayout>(op)) { in ExtractLayoutFromOperand() local
Dvalue_utils.cc49 if (mlir::TF::DTensorLayout layout_op = in GetForwardedInput() local
Dhandle_cross_cluster_dependencies.cc99 auto layout_op = builder.create<mlir::TF::DTensorLayout>( in CloneOpToCluster() local
Dmesh_propagation.cc208 if (auto layout_op = in InferMeshFromInputs() local
Dspmd_expander_common.cc322 auto layout_op = in GetForwardedDTensorLayoutInput() local
Dlayout_propagation_v2.cc699 auto layout_op = builder.create<mlir::TF::DTensorLayout>( in InsertDTensorLayoutOps() local
1024 for (auto layout_op : llvm::reverse(layout_ops)) { in ReplaceAuxiliaryDTensorLayoutOpsWithIdentity() local