Searched defs:all_scatter (Results 1 – 3 of 3) sorted by relevance
50 mlir::TF::DTensorAllScatterOp all_scatter, int scatter_dim) { in GetScatterGroupAssignment()81 if (auto all_scatter = mlir::dyn_cast<mlir::TF::DTensorAllScatterOp>( in ApplyOptimization() local
137 mlir::TF::DTensorAllScatterOp all_scatter = in EmitAllScatter() local274 auto all_scatter = in EmitRelayout() local
630 mlir::TF::DTensorAllScatterOp all_scatter) { in LowerAllScatterOp()795 module.walk([&](mlir::TF::DTensorAllScatterOp all_scatter) { in runOnOperation()799 for (mlir::TF::DTensorAllScatterOp all_scatter : all_scatters) in runOnOperation() local