# Path: smore/common/embedding/sparse_embed.py
# class SparseEmbedding(nn.Module):
#     """Sparse embedding"""
# 
#     def __init__(self, num_rows, num_cols):
#         super(SparseEmbedding, self).__init__()
#         self.num_rows = num_rows
#         self.embedding = nn.Parameter(torch.zeros(num_rows, num_cols))
#         self.fwd_delegate = None
#         self.dummy_job = DummyJob()
# 
#     @property
#     def device(self):
#         return self.embedding.device
# 
#     @property
#     def shape(self):
#         return self.embedding.shape
# 
#     def init_params(self, lb, ub):
#         nn.init.uniform_(
#             tensor=self.embedding, 
#             a=lb,
#             b=ub
#         )
# 
#     def register_tensor(self, name, tensor_obj):
#         self.register_buffer(name, tensor_obj)
#         if self.embedding.is_shared():
#             tensor_obj.share_memory_()
# 
#     def forward(self, indices, name=None):
#         if self.num_rows == 0:
#             return None
#         if self.fwd_delegate is not None:
#             return self.fwd_delegate(indices, name)
#         else:
#             if indices is None:
#                 mat = self.embedding
#             else:
#                 mat = self.embedding[indices]
#             mat.job_handle = self.dummy_job
#             return mat
# 
# Path: smore/common/embedding/embed_rw.py
# class EmbeddingRW(EmbeddingReadOnly):
#     def __init__(self, embed, gpu_id=-1):
#         super(EmbeddingRW, self).__init__(embed, gpu_id)
#         self.write_thread_pool = self.read_thread_pool
#         self.write_buf = {}        
#         self.write_src_cache = {}
# 
#     def write(self, indices, values, name=None, additive=False):
#         if self.embed.is_cuda:
#             indices = indices.to(self.embed.device)
#         elif name is not None:
#             return self.async_write(indices, values, name, additive)
#         if additive:
#             self.embed.index_add_(0, indices, values.to(self.embed.device))
#         else:
#             self.embed.index_copy_(0, indices, values.to(self.embed.device))
#         return self.dummy_job
# 
#     def async_write(self, indices, values, name, additive=False):
#         src_event = extlib.CudaEvent(self.gpu_id)
#         src_event.record()
#         indices = indices.view(-1)
#         assert indices.is_cuda == False and values.is_cuda
#         assert values.shape[0] == indices.shape[0] and values.shape[1] == self.embed_dim
#         if name in self.last_write_jobs:
#             self.last_write_jobs[name].sync()
#         buf, _ = self.get_or_create_buf(name, indices.shape[0], self.write_buf, gpu_buf=False)
#         if buf.shape[0] > indices.shape[0]:
#             buf = buf[:indices.shape[0]]
#         with torch.cuda.stream(self.stream):
#             job_handle = extlib.async_write(self.write_thread_pool, 
#                                             indices,
#                                             self.embed,
#                                             buf,
#                                             values,
#                                             src_event,
#                                             additive)
#         self.write_src_cache[name] = (src_event, indices, values)
#         self.last_write_jobs[name] = job_handle
#         return job_handle
#  
#     def add(self, indices, values, name="default"):
#         return self.write(indices, values, name, additive=True)
# 
# Path: smore/common/util.py
# def eval_tuple(arg_return):
#     """Evaluate a tuple string into a tuple."""
#     if type(arg_return) == tuple:
#         return arg_return
#     if arg_return[0] not in ["(", "["]:
#         arg_return = eval(arg_return)
#     else:
#         splitted = arg_return[1:-1].split(",")
#         List = []
#         for item in splitted:
#             try:
#                 item = eval(item)
#             except:
#                 pass
#             if item == "":
#                 continue
#             List.append(item)
#         arg_return = tuple(List)
#     return arg_return
# 
# Path: smore/common/consts.py



import time
import math
import numpy as np
import torch
import torch.multiprocessing as mp
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable
from torch.nn.parameter import Parameter
from torch.multiprocessing import Queue, Pipe
from torch_scatter import scatter
from functools import partial
from abc import ABC, abstractmethod
from smore.common.embedding.sparse_embed import SparseEmbedding
from smore.common.embedding.embed_rw import EmbeddingRW
from smore.common.util import eval_tuple
from smore.common.consts import eps
        squeeze = optim_mode[3]
        if squeeze:
            fn_register("state_sum", torch.zeros(embed.embedding.size(0), 1))
        else:
            fn_register("state_sum", torch.zeros(embed.embedding.size()))

    def share_memory(self):
        self.state_sum_rw.embed.share_memory_()

    def prepare_forward(self, indices, name):
        fo = self.state_sum_rw.read(indices, name)
        self.grad_stats.append(fo)

    def update_step(self, indices, inverse_indices, grad, lr):
        # indices, grad on gpu
        stats_device = self.state_sum_rw.embed.device
        with torch.no_grad():
            if self.optimizer_device == 'gpu':
                indices = indices.to(stats_device)
            if self.squeeze:
                grad_sum = torch.mean(grad * grad, dim=1, keepdim=True) # gpu
            else:
                grad_sum = grad * grad

            if self.optimizer_device == 'cpu':
                list_sq = self.filter_stats(self.grad_stats)
                square_avg = merge_grad_list(list_sq, inverse_indices, reduce_type="max")
            else:
                square_avg = self.state_sum_rw.read(indices)
            square_avg = square_avg + grad_sum

# Path: smore/common/embedding/sparse_embed.py
# class SparseEmbedding(nn.Module):
#     """Sparse embedding"""
# 
#     def __init__(self, num_rows, num_cols):
#         super(SparseEmbedding, self).__init__()
#         self.num_rows = num_rows
#         self.embedding = nn.Parameter(torch.zeros(num_rows, num_cols))
#         self.fwd_delegate = None
#         self.dummy_job = DummyJob()
# 
#     @property
#     def device(self):
#         return self.embedding.device
# 
#     @property
#     def shape(self):
#         return self.embedding.shape
# 
#     def init_params(self, lb, ub):
#         nn.init.uniform_(
#             tensor=self.embedding, 
#             a=lb,
#             b=ub
#         )
# 
#     def register_tensor(self, name, tensor_obj):
#         self.register_buffer(name, tensor_obj)
#         if self.embedding.is_shared():
#             tensor_obj.share_memory_()
# 
#     def forward(self, indices, name=None):
#         if self.num_rows == 0:
#             return None
#         if self.fwd_delegate is not None:
#             return self.fwd_delegate(indices, name)
#         else:
#             if indices is None:
#                 mat = self.embedding
#             else:
#                 mat = self.embedding[indices]
#             mat.job_handle = self.dummy_job
#             return mat
# 
# Path: smore/common/embedding/embed_rw.py
# class EmbeddingRW(EmbeddingReadOnly):
#     def __init__(self, embed, gpu_id=-1):
#         super(EmbeddingRW, self).__init__(embed, gpu_id)
#         self.write_thread_pool = self.read_thread_pool
#         self.write_buf = {}        
#         self.write_src_cache = {}
# 
#     def write(self, indices, values, name=None, additive=False):
#         if self.embed.is_cuda:
#             indices = indices.to(self.embed.device)
#         elif name is not None:
#             return self.async_write(indices, values, name, additive)
#         if additive:
#             self.embed.index_add_(0, indices, values.to(self.embed.device))
#         else:
#             self.embed.index_copy_(0, indices, values.to(self.embed.device))
#         return self.dummy_job
# 
#     def async_write(self, indices, values, name, additive=False):
#         src_event = extlib.CudaEvent(self.gpu_id)
#         src_event.record()
#         indices = indices.view(-1)
#         assert indices.is_cuda == False and values.is_cuda
#         assert values.shape[0] == indices.shape[0] and values.shape[1] == self.embed_dim
#         if name in self.last_write_jobs:
#             self.last_write_jobs[name].sync()
#         buf, _ = self.get_or_create_buf(name, indices.shape[0], self.write_buf, gpu_buf=False)
#         if buf.shape[0] > indices.shape[0]:
#             buf = buf[:indices.shape[0]]
#         with torch.cuda.stream(self.stream):
#             job_handle = extlib.async_write(self.write_thread_pool, 
#                                             indices,
#                                             self.embed,
#                                             buf,
#                                             values,
#                                             src_event,
#                                             additive)
#         self.write_src_cache[name] = (src_event, indices, values)
#         self.last_write_jobs[name] = job_handle
#         return job_handle
#  
#     def add(self, indices, values, name="default"):
#         return self.write(indices, values, name, additive=True)
# 
# Path: smore/common/util.py
# def eval_tuple(arg_return):
#     """Evaluate a tuple string into a tuple."""
#     if type(arg_return) == tuple:
#         return arg_return
#     if arg_return[0] not in ["(", "["]:
#         arg_return = eval(arg_return)
#     else:
#         splitted = arg_return[1:-1].split(",")
#         List = []
#         for item in splitted:
#             try:
#                 item = eval(item)
#             except:
#                 pass
#             if item == "":
#                 continue
#             List.append(item)
#         arg_return = tuple(List)
#     return arg_return
# 
# Path: smore/common/consts.py



# Path: smore/common/embedding/embed_optimizer.py
import time
import math
import numpy as np
import torch
import torch.multiprocessing as mp
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable
from torch.nn.parameter import Parameter
from torch.multiprocessing import Queue, Pipe
from torch_scatter import scatter
from functools import partial
from abc import ABC, abstractmethod
from smore.common.embedding.sparse_embed import SparseEmbedding
from smore.common.embedding.embed_rw import EmbeddingRW
from smore.common.util import eval_tuple
from smore.common.consts import eps

        squeeze = optim_mode[3]
        if squeeze:
            fn_register("state_sum", torch.zeros(embed.embedding.size(0), 1))
        else:
            fn_register("state_sum", torch.zeros(embed.embedding.size()))

    def share_memory(self):
        self.state_sum_rw.embed.share_memory_()

    def prepare_forward(self, indices, name):
        fo = self.state_sum_rw.read(indices, name)
        self.grad_stats.append(fo)

    def update_step(self, indices, inverse_indices, grad, lr):
        # indices, grad on gpu
        stats_device = self.state_sum_rw.embed.device
        with torch.no_grad():
            if self.optimizer_device == 'gpu':
                indices = indices.to(stats_device)
            if self.squeeze:
                grad_sum = torch.mean(grad * grad, dim=1, keepdim=True) # gpu
            else:
                grad_sum = grad * grad

            if self.optimizer_device == 'cpu':
                list_sq = self.filter_stats(self.grad_stats)
                square_avg = merge_grad_list(list_sq, inverse_indices, reduce_type="max")
            else:
                square_avg = self.state_sum_rw.read(indices)
            square_avg = square_avg + grad_sum

            avg = square_avg.sqrt().add_(eps)