def update(self, idx, grad, emb):
        """Update embeddings in a sparse manner
        Sparse embeddings are updated in mini batches. We maintain gradient states for
        each embedding so they can be updated separately.

        Parameters
        ----------
        idx : tensor
            Index of the embeddings to be updated.
        grad : tensor
            Gradient of each embedding.
        emb : dgl.nn.NodeEmbedding
            Sparse embedding to update.
        """
        with th.no_grad():
            state_step, state_mem, state_power = emb.optm_state
            exec_dtype = grad.dtype
            exec_dev = grad.device
            state_dev = state_step.device

            # whether or not we need to transfer data from the GPU to the CPU
            # while updating the weights
            is_d2h = state_dev.type == "cpu" and exec_dev.type == "cuda"

            # only perform async copies cpu -> gpu, or gpu-> gpu, but block
            # when copying to the cpu, so as to ensure the copy is finished
            # before operating on the data on the cpu
            state_block = is_d2h

            if self._is_using_uva[emb.name] is None and is_d2h:
                # we should use UVA going forward
                self._setup_uva(emb.name, state_mem, state_power)
            elif self._is_using_uva[emb.name] is None:
                # we shouldn't use UVA going forward
                self._is_using_uva[emb.name] = False

            use_uva = self._is_using_uva[emb.name]

            beta1 = self._beta1
            beta2 = self._beta2
            eps = self._eps

            clr = self._lr
            # There can be duplicated indices due to sampling.
            # Thus unique them here and average the gradient here.
            grad_indices, inverse, cnt = th.unique(
                idx, return_inverse=True, return_counts=True
            )
            state_idx = grad_indices.to(state_dev)
            state_step[state_idx] += 1
            state_step = state_step[state_idx].to(exec_dev)

            if use_uva:
                orig_mem = gather_pinned_tensor_rows(state_mem, grad_indices)
                orig_power = gather_pinned_tensor_rows(
                    state_power, grad_indices
                )
            else:
                orig_mem = state_mem[state_idx].to(exec_dev)
                orig_power = state_power[state_idx].to(exec_dev)
            # convert to exec dtype
            orig_mem = orig_mem.to(dtype=exec_dtype)
            orig_power = orig_power.to(dtype=exec_dtype)

            grad_values = th.zeros(
                (grad_indices.shape[0], grad.shape[1]), device=exec_dev
            )
            grad_values.index_add_(0, inverse, grad)
            grad_values = grad_values / cnt.unsqueeze(1)

            grad_mem = grad_values
            grad_power = grad_values * grad_values

            update_mem = beta1 * orig_mem + (1.0 - beta1) * grad_mem
            update_power = beta2 * orig_power + (1.0 - beta2) * grad_power

            if use_uva:
                scatter_pinned_tensor_rows(
                    state_mem, grad_indices, update_mem.to(dtype=self._dtype)
                )
                scatter_pinned_tensor_rows(
                    state_power,
                    grad_indices,
                    update_power.to(dtype=self._dtype),
                )
            else:
                update_mem_dst = update_mem.to(dtype=self._dtype).to(
                    state_dev, non_blocking=True
                )
                update_power_dst = update_power.to(dtype=self._dtype).to(
                    state_dev, non_blocking=True
                )
                if state_block:
                    # use events to try and overlap CPU and GPU as much as possible
                    update_event = th.cuda.Event()
                    update_event.record()

            update_mem_corr = update_mem / (
                1.0 - th.pow(th.tensor(beta1, device=exec_dev), state_step)
            ).unsqueeze(1)
            update_power_corr = update_power / (
                1.0 - th.pow(th.tensor(beta2, device=exec_dev), state_step)
            ).unsqueeze(1)
            std_values = (
                clr * update_mem_corr / (th.sqrt(update_power_corr) + eps)
            )
            std_values_dst = std_values.to(state_dev, non_blocking=True)

            if state_block:
                std_event = th.cuda.Event()
                std_event.record()

            if not use_uva:
                if state_block:
                    # wait for our transfers from exec_dev to state_dev to finish
                    # before we can use them
                    update_event.wait()
                state_mem[state_idx] = update_mem_dst
                state_power[state_idx] = update_power_dst

            if state_block:
                # wait for the transfer of std_values to finish before we
                # can use it
                std_event.wait()
            emb.weight[state_idx] -= std_values_dst
Update embeddings in a sparse manner
        Sparse embeddings are updated in mini batches. We maintain gradient states for
        each embedding so they can be updated separately.

        Parameters
        ----------
        idx : tensor
            Index of the embeddings to be updated.
        grad : tensor
            Gradient of each embedding.
        emb : dgl.nn.NodeEmbedding
            Sparse embedding to update.
        
python/dgl/optim/pytorch/sparse_optim.py