[GraphBolt] move return_eids check to internal python API (#7071)

python/dgl/graphbolt/impl/fused_csc_sampling_graph.py

@@ -625,8 +625,16 @@ class FusedCSCSamplingGraph(SamplingGraph):
         if isinstance(nodes, dict):
             nodes = self._convert_to_homogeneous_nodes(nodes)
 
+        return_eids = (
+            self.edge_attributes is not None
+            and ORIGINAL_EDGE_ID in self.edge_attributes
+        )
         C_sampled_subgraph = self._sample_neighbors(
-            nodes, fanouts, replace, probs_name
+            nodes,
+            fanouts,
+            replace=replace,
+            probs_name=probs_name,
+            return_eids=return_eids,
         )
         return self._convert_to_sampled_subgraph(C_sampled_subgraph)
 
@@ -679,6 +687,7 @@ class FusedCSCSamplingGraph(SamplingGraph):
         fanouts: torch.Tensor,
         replace: bool = False,
         probs_name: Optional[str] = None,
+        return_eids: bool = False,
     ) -> torch.ScriptObject:
         """Sample neighboring edges of the given nodes and return the induced
         subgraph.
@@ -714,6 +723,9 @@ class FusedCSCSamplingGraph(SamplingGraph):
             corresponding to each neighboring edge of a node. It must be a 1D
             floating-point or boolean tensor, with the number of elements
             equalling the total number of edges.
+        return_eids: bool, optional
+            Boolean indicating whether to return the original edge IDs of the
+            sampled edges.
 
         Returns
         -------
@@ -722,16 +734,12 @@ class FusedCSCSamplingGraph(SamplingGraph):
         """
         # Ensure nodes is 1-D tensor.
         self._check_sampler_arguments(nodes, fanouts, probs_name)
-        has_original_eids = (
-            self.edge_attributes is not None
-            and ORIGINAL_EDGE_ID in self.edge_attributes
-        )
         return self._c_csc_graph.sample_neighbors(
             nodes,
             fanouts.tolist(),
             replace,
             False,
-            has_original_eids,
+            return_eids,
             probs_name,
         )