def _set_e_repr(self, etid, edges, data):
        """Internal API to set edge(s) features.

        `data` is a dictionary from the feature name to feature tensor. Each tensor
        is of shape (B, D1, D2, ...), where B is the number of edges to be updated,
        and (D1, D2, ...) be the shape of the edge representation tensor.

        All update will be done out of place to work with autograd.

        Parameters
        ----------
        etid : int
            Edge type id.
        edges : edges
            Edges can be either

            * A pair of endpoint nodes (u, v), where u is the node ID of source
              node type and v is that of destination node type.
            * A tensor of edge ids of the given type.

            The default value is all the edges.
        data : tensor or dict of tensor
            Edge representation.
        """
        # parse argument
        if not is_all(edges):
            eid = utils.parse_edges_arg_to_eid(self, edges, etid, "edges")

        # sanity check
        if not utils.is_dict_like(data):
            raise DGLError(
                "Expect dictionary type for feature data."
                ' Got "%s" instead.' % type(data)
            )

        if is_all(edges):
            num_edges = self._graph.num_edges(etid)
        else:
            num_edges = len(eid)
        for key, val in data.items():
            nfeats = F.shape(val)[0]
            if nfeats != num_edges:
                raise DGLError(
                    "Expect number of features to match number of edges."
                    " Got %d and %d instead." % (nfeats, num_edges)
                )
            if F.context(val) != self.device:
                raise DGLError(
                    'Cannot assign edge feature "{}" on device {} to a graph on'
                    " device {}. Call DGLGraph.to() to copy the graph to the"
                    " same device.".format(key, F.context(val), self.device)
                )
            # To prevent users from doing things like:
            #
            #     g.pin_memory_()
            #     g.edata['x'] = torch.randn(...)
            #     sg = g.sample_neighbors(torch.LongTensor([...]).cuda())
            #     sg.edata['x']    # Becomes a CPU tensor even if sg is on GPU due to lazy slicing
            if (
                self.is_pinned()
                and F.context(val) == "cpu"
                and not F.is_pinned(val)
            ):
                raise DGLError(
                    "Pinned graph requires the edge data to be pinned as well. "
                    "Please pin the edge data before assignment."
                )

        # set
        if is_all(edges):
            self._edge_frames[etid].update(data)
        else:
            self._edge_frames[etid].update_row(eid, data)
Internal API to set edge(s) features.

        `data` is a dictionary from the feature name to feature tensor. Each tensor
        is of shape (B, D1, D2, ...), where B is the number of edges to be updated,
        and (D1, D2, ...) be the shape of the edge representation tensor.

        All update will be done out of place to work with autograd.

        Parameters
        ----------
        etid : int
            Edge type id.
        edges : edges
            Edges can be either

            * A pair of endpoint nodes (u, v), where u is the node ID of source
              node type and v is that of destination node type.
            * A tensor of edge ids of the given type.

            The default value is all the edges.
        data : tensor or dict of tensor
            Edge representation.
        
python/dgl/heterograph.py