if (all_uses_replaced) {
      op.erase();
    }
    return WalkResult::advance();
  });
}

}  // namespace

std::unique_ptr<OperationPass<ModuleOp>>
CreateConstantOpDeviceAssignmentPass() {
  return std::make_unique<ConstantOpDeviceAssignmentPass>();
}

}  // namespace TF

// a tf.Const op is read by operations placed on multiple devices, then the pass
// will replicate the tf.Const op once for each device.
std::unique_ptr<OperationPass<ModuleOp>> CreateConstantOpDeviceAssignmentPass();

// Returns pass that verifies whether all functions in module are of single
// tf_executor.graph and each tf_executor.island in tf_executor.graph only has a
// single op.

  ];
}

def ConstantOpDeviceAssignmentPass : Pass<"constant-op-device-assignment", "ModuleOp"> {
  let summary = "Assign device for tf.Const ops";
  let constructor = "TF::CreateConstantOpDeviceAssignmentPass()";
}

def GuaranteeAllFuncsOneUsePass : Pass<"tf-guarantee-all-funcs-one-use", "ModuleOp"> {
  let summary = "Guarantee all FuncOp's have only a single use.";