terpret_problem.py

import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable, Function
from torch.optim.optimizer import Optimizer

from typing import TypeVar, Union, Tuple, Optional, Callable

import numpy as np
import sys


class TerpretProblem(nn.Module):
    def __init__(self,opts):
        super(TerpretProblem, self).__init__()
        self.opts=opts
        self.ms = nn.Parameter(torch.rand(opts.k-1, 2,dtype=torch.float), requires_grad=True)

    def forward(self, x):
        opts=self.opts
        self.mus = torch.cat([x, torch.nn.functional.softmax(self.ms,dim=1)], 0)
        def soft_xor_p0(i):
            j = i+1
            j = j % opts.k
            return self.mus[i, 0] * self.mus[j, 0] + self.mus[i, 1] * self.mus[j, 1]

        self.ys_eq_0 = torch.stack([ soft_xor_p0(temp) for temp in range(opts.k)],0)

        self.result = torch.log(self.ys_eq_0).sum()
        return - self.result

    def return_0_prob(self,x):
        return self.mus


class TerpretProblem_ConcreteDistribution(TerpretProblem):
    def __init__(self,opts):
        super(TerpretProblem_ConcreteDistribution, self).__init__(opts)
        self.initial_temp = 2.5
         # innitial gumbel-softmax temperature
        self.temperature = self.initial_temp
        # annealation rate of softmax temperature

        self.anneal_rate = 0.9999
        self.anneal_rate = 0.999
        self.min_temperature = 0.1

        self.step=0
        self.M=opts.M # the number of samples

    def forward(self, x):
        self.step+=1
        # self.temperature = max(self.temperature * self.anneal_rate,self.min_temperature)
        self.temperature = self.temperature * self.anneal_rate
        def gumbel_softmax_sample(logits, temperature):
            m = torch.distributions.relaxed_categorical.RelaxedOneHotCategorical(
                torch.tensor([temperature]), logits=logits
                )
            # return m.rsample(),-m.log_prob(m.sample())
            return m.rsample() # note that `rsample` is reparameterization.

        opts=self.opts

        self.result=0
        for i in range(self.M):
            self.ms_sampled = gumbel_softmax_sample(self.ms,self.temperature)

            self.mus = torch.cat([x, self.ms_sampled], 0)
            def soft_xor_p0(i):
                j = i+1
                j = j % opts.k
                return self.mus[i, 0] * self.mus[j, 0] + self.mus[i, 1] * self.mus[j, 1]

            self.ys_eq_0 = torch.stack([ soft_xor_p0(temp) for temp in range(opts.k)],0)
            self.result += torch.log(self.ys_eq_0).sum()

        return - self.result / self.M

    def return_0_prob(self,x):
        self.mus = torch.cat([x, torch.nn.functional.softmax(self.ms,dim=1)], 0)
        return self.mus

class Binarize(Function):
    '''used where binarization of real-valued parameter into binary values are needed.'''
    clip_value = 0.5

    @staticmethod
    def forward(ctx, inp):
        ctx.save_for_backward(inp)

        # output = inp.sign()
        output = torch.where(inp > Binarize.clip_value, torch.ones_like(inp).type_as(inp), torch.zeros_like(inp).type_as(inp))
        return output

    @staticmethod
    def backward(ctx, grad_output):
        inp: Tensor = ctx.saved_tensors[0]

        # clipped = inp.abs() <= Binarize.clip_value
        clipped = (inp-0.5).abs() <= 0.5

        output = torch.zeros(inp.size()).to(grad_output.device)
        output[clipped] = 1
        output[~clipped] = 0

        return output * grad_output

binarize = Binarize.apply


class CustomizedAdam(Optimizer):
    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8):
        if not 0.0 <= lr:
            raise ValueError("Invalid learning rate: {}".format(lr))
        if not 0.0 <= eps:
            raise ValueError("Invalid epsilon value: {}".format(eps))
        if not 0.0 <= betas[0] < 1.0:
            raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
        if not 0.0 <= betas[1] < 1.0:
            raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))

        defaults = dict(lr=lr, betas=betas, eps=eps,)
        super(CustomizedAdam, self).__init__(params, defaults)

    def step(self, closure: Optional[Callable[[], float]] = ...):
        for group in self.param_groups:
            for p in group['params']:
                if p.grad is not None:
                    params = p
                    if p.grad.is_sparse:
                        raise RuntimeError('Adam does not support sparse gradients, please consider SparseAdam instead')
                    grad = p.grad.data

                    state = self.state[p]
                    # Lazy state initialization
                    if len(state) == 0:
                        state['step'] = 0
                        # Exponential moving average of gradient values
                        state['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)
                        # Exponential moving average of squared gradient values
                        state['exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)

                    exp_avg = state['exp_avg']
                    exp_avg_sq = state['exp_avg_sq']

                    # update the steps for each param group update
                    state['step'] += 1
                    # record the step after step update
                    state_steps = state['step']

                    beta1,beta2 = group["betas"]
                    step_size = group["lr"]

                    # Decay the first and second moment running average coefficient
                    exp_avg.mul_(beta1).add_(1 - beta1, grad)
                    exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)

                    denom = exp_avg_sq.sqrt().add_(group['eps'])

                    bias_correction1 = 1 - beta1 ** state['step']
                    bias_correction2 = 1 - beta2 ** state['step']
                    step_size = group['lr'] * math.sqrt(bias_correction2) / bias_correction1

                    p.data.addcdiv_(-step_size, exp_avg, denom)

        return

    def zero_grad(self) -> None:
        super().zero_grad()
        if self.use_non_binary is True:
            self._adam.zero_grad()


class NewAdam(Optimizer):
    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8):
        if not 0.0 <= lr:
            raise ValueError("Invalid learning rate: {}".format(lr))
        if not 0.0 <= eps:
            raise ValueError("Invalid epsilon value: {}".format(eps))
        if not 0.0 <= betas[0] < 1.0:
            raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
        if not 0.0 <= betas[1] < 1.0:
            raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))

        defaults = dict(lr=lr, betas=betas, eps=eps,)
        super(NewAdam, self).__init__(params, defaults)

    def step(self, closure: Optional[Callable[[], float]] = ...):
        for group in self.param_groups:
            for p in group['params']:
                if p.grad is not None:
                    params = p
                    if p.grad.is_sparse:
                        raise RuntimeError('Adam does not support sparse gradients, please consider SparseAdam instead')
                    grad = p.grad.data

                    state = self.state[p]
                    # Lazy state initialization
                    if len(state) == 0:
                        # state['step'] = 0
                        state['step'] = torch.zeros_like(p, memory_format=torch.preserve_format)
                        # Exponential moving average of gradient values
                        state['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)
                        # Exponential moving average of squared gradient values
                        state['exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)

                    exp_avg = state['exp_avg']
                    exp_avg_sq = state['exp_avg_sq']

                    # update the steps for each param group update
                    state['step'] += 1
                    # record the step after step update
                    state_steps = state['step']

                    beta1,beta2 = group["betas"]

                    # Decay the first and second moment running average coefficient
                    exp_avg.mul_(beta1).add_(1 - beta1, grad)
                    exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)

                    denom = exp_avg_sq.sqrt().add_(group['eps'])

                    bias_correction1 = 1 - beta1 ** state['step']
                    # bias_correction1 = 1 - torch.zeros_like(p)
                    bias_correction2 = 1 - beta2 ** state['step']
                    bias_correction2 = 1 - torch.zeros_like(p)

                    step_size = group['lr'] * torch.sqrt(bias_correction2) / bias_correction1
                    # step_size = group['lr']

                    # record the sign before update
                    sign_old = (p.data > 0.5 )

                    p.data.add_(-step_size* (exp_avg/denom) )
                    # p.data.add_(-step_size* exp_avg )

                    # record the sign after update
                    sign_new = (p.data > 0.5)

                    # reset the moment according to the change of sign
                    reset_mask = ( sign_new == sign_old).float()
                    # now the entries with a change of sign is zero, and the ones without change is one.
                    state_steps.mul_(reset_mask)
                    exp_avg.mul_(reset_mask)
                    # exp_avg_sq.mul_(1-reset_mask)



        return

    def zero_grad(self) -> None:
        super().zero_grad()
        if self.use_non_binary is True:
            self._adam.zero_grad()



class TerpretProblem_STE(TerpretProblem):
    def __init__(self,opts,adaptive_noise=False):
        super(TerpretProblem_STE, self).__init__(opts)
        self.opts=opts
        self.parameter = nn.Parameter(torch.rand(opts.k-1, 1,dtype=torch.float)*0.1+0.45, requires_grad=True)
        # self.parameter = nn.Parameter(torch.rand(opts.k-1, 1,dtype=torch.float), requires_grad=True)



    def forward(self, x):
        opts=self.opts
        ms = torch.cat((1-F.sigmoid(self.parameter-0.5), F.sigmoid(self.parameter-0.5)), dim=1)
        self.mus_latent = torch.cat([x, ms], 0)

        temp = binarize(self.parameter)
        ms = torch.cat((1-temp, temp), dim=1)
        self.mus = torch.cat([x, ms], 0)

        def soft_xor_p0(i):
            j = i+1
            j = j % opts.k
            return self.mus[i, 0] * self.mus[j, 0] + self.mus[i, 1] * self.mus[j, 1]

        self.ys_eq_0 = torch.stack([ soft_xor_p0(temp) for temp in range(opts.k)],0)

        eps=1e-7
        self.result = torch.log(self.ys_eq_0+eps).sum()
        return - self.result

class Bop(Optimizer):
    def __init__(
        self,
        binary_params,
        ar = 0.0001,
        threshold = 0.00001,
        continuous_optimizer = None
    ):
        if not 0 < ar < 1:
            raise ValueError(
                "given adaptivity rate {} is invalid; should be in (0, 1) (excluding endpoints)".format(
                    ar
                )
            )
        if threshold < 0:
            raise ValueError(
                "given threshold {} is invalid; should be > 0".format(threshold)
            )

        self.total_weights = {}
        if continuous_optimizer is None:
            self.use_non_binary=False
        else:
            self.use_non_binary=True
            self._adam = continuous_optimizer


        defaults = dict(adaptivity_rate=ar, threshold=threshold)
        super(Bop, self).__init__(
            binary_params, defaults
        )

    def step(self, closure: Optional[Callable[[], float]] = ..., ar=None):
        if self.use_non_binary is True:
            self._adam.step()

        flips = {None}

        for group in self.param_groups:
            params = group["params"]

            y = group["adaptivity_rate"]
            t = group["threshold"]
            flips = {}

            if ar is not None:
                y = ar

            for param_idx, p in enumerate(params):
                grad = p.grad.data

                state = self.state[p]
                if len(state) == 0:
                    state['step'] = 0
                    state["moving_average"] = torch.zeros_like(p.data, memory_format=torch.preserve_format)
                    # state["moving_average"] = torch.zeros_like(torch.rand(p.data.size()), memory_format=torch.preserve_format)

                m = state['moving_average']
                m.mul_((1 - y))
                m.add_(grad.mul(y))
                # print(m)
                m_=m

                p_temp = torch.where(p>0.5, torch.ones_like(p), - torch.ones_like(p))

                mask = (m_.abs() >= t ) * (m_.sign() == p_temp.sign())
                mask = mask.float() * 1
                flips[param_idx] = mask

                p.data.logical_xor_(mask).type(torch.float)

            # print('# of flip',flips)
        return flips,grad

    def zero_grad(self) -> None:
        super().zero_grad()
        if self.use_non_binary is True:
            self._adam.zero_grad()

class TerpretProblem_Bop(TerpretProblem):
    def __init__(self,opts,use_adaptive_noise=False):
        super(TerpretProblem_Bop, self).__init__(opts)
        self.opts=opts

        self.use_adaptive_noise=use_adaptive_noise

        random_noise = torch.rand(opts.k-1, 1,dtype=torch.float)
        init_value = torch.where(random_noise>0.5,torch.ones_like(random_noise),torch.zeros_like(random_noise))
        self.parameter = nn.Parameter(init_value, requires_grad=True)

        if self.use_adaptive_noise is True:
            noise_init = 0.00000001
            noise_init = 0.0

            self.noise_rate = nn.Parameter(noise_init*torch.ones(opts.k-1,1))

    def binary_parameters(self):
        return [self.parameter]

    def continuous_parameters(self):
        return [self.noise_rate]

    def forward(self, x):
        opts=self.opts
        temp = self.parameter

        if self.use_adaptive_noise is True:
            noise = torch.where(self.parameter > 0.5, -torch.rand(self.parameter.shape).type_as(self.parameter), torch.rand(self.parameter.shape).type_as(self.parameter))
            temp = self.parameter + (self.noise_rate.abs() ) * noise

        ms = torch.cat((1-temp, temp), dim=1)
        self.mus = torch.cat([x, ms], 0)

        def soft_xor_p0(i):
            j = i+1
            j = j % opts.k
            return self.mus[i, 0] * self.mus[j, 0] + self.mus[i, 1] * self.mus[j, 1]


        self.ys_eq_0 = torch.stack([ soft_xor_p0(temp) for temp in range(opts.k)],0)

        eps=1e-7
        self.result = torch.log(self.ys_eq_0+eps).sum()
        return - self.result

class TerpretProblem_binary(nn.Module):
    def __init__(self,opts):
        super(TerpretProblem_binary, self).__init__()
        self.opts=opts


        random_noise = torch.rand(opts.k-1, 1,dtype=torch.float)
        init_value = torch.where(random_noise>0.5,torch.ones_like(random_noise),torch.zeros_like(random_noise))
        # init_value = torch.cat((1-init_value, init_value), dim=1)
        self.parameter = nn.Parameter(init_value, requires_grad=True)

        self.ms = self.parameter

    def binary_parameters(self):
        return [self.parameter]

    def forward(self, x):
        opts=self.opts
        # self.mus = torch.cat([x, self.parameter], 0)

        temp = self.parameter
        ms = torch.cat((1-temp, temp), dim=1)
        self.mus = torch.cat([x, ms], 0)


        def soft_xor_p0(i):
            j = i+1
            j = j % opts.k
            return self.mus[i, 0] * self.mus[j, 0] + self.mus[i, 1] * self.mus[j, 1]

        self.ys_eq_0 = torch.stack([ soft_xor_p0(temp) for temp in range(opts.k)],0)

        eps=1e-7
        self.result = torch.log(self.ys_eq_0+eps).sum()
        return - self.result

    def return_0_prob(self,x):
        return self.mus


class Reparameterization_Net(nn.Module):
    def __init__(self,size):
        super(Reparameterization_Net, self).__init__()
        self.output_size = size # the size of the output shape

        self.output_dimension = self.get_output_flattened_size()
        self.input_dimension = self.get_input_flattened_size()
        self.intermediate_dimension = self.output_dimension

        self.computation = nn.Sequential(
            nn.Linear(self.input_dimension,self.intermediate_dimension),
            nn.ReLU(),
            nn.Linear(self.intermediate_dimension,self.intermediate_dimension),
            nn.ReLU(),
            nn.Linear(self.intermediate_dimension,self.intermediate_dimension),
            nn.ReLU(),
            nn.Linear(self.intermediate_dimension,self.intermediate_dimension),
            nn.ReLU(),
            nn.Linear(self.intermediate_dimension,self.intermediate_dimension),
            nn.ReLU(),
            nn.Linear(self.intermediate_dimension,self.intermediate_dimension),
            nn.ReLU(),
            nn.Linear(self.intermediate_dimension,self.intermediate_dimension),
            nn.ReLU(),
            nn.Linear(self.intermediate_dimension,self.intermediate_dimension),
            nn.ReLU(),
            nn.Linear(self.intermediate_dimension,self.intermediate_dimension),
            nn.ReLU(),
            nn.Linear(self.intermediate_dimension,self.intermediate_dimension),
            nn.ReLU(),
            nn.Linear(self.intermediate_dimension,self.intermediate_dimension),
            nn.ReLU(),
            nn.Linear(self.intermediate_dimension,self.intermediate_dimension),
            nn.ReLU(),
            nn.Linear(self.intermediate_dimension,self.intermediate_dimension),
            nn.ReLU(),
            nn.Linear(self.intermediate_dimension,self.intermediate_dimension),
            nn.ReLU(),
            nn.Linear(self.intermediate_dimension,self.intermediate_dimension),
            nn.ReLU(),
            nn.Linear(self.intermediate_dimension,self.output_dimension),

            #
            # nn.Linear(self.input_dimension,int(self.input_dimension/10)),
            # nn.ReLU(),
            # nn.Linear(int(self.input_dimension/10),int(self.input_dimension/10)),
            # nn.ReLU(),
            # nn.Linear(int(self.input_dimension/10),int(self.input_dimension/100)),
            # nn.ReLU(),
            # nn.Linear(int(self.input_dimension/100),int(self.input_dimension/10)),
            # nn.ReLU(),
            # nn.Linear(int(self.input_dimension/10),int(self.input_dimension/10)),
            # nn.ReLU(),
            # nn.Linear(self.intermediate_dimension,self.output_dimension),
            # nn.ReLU()



        )

    def get_output_flattened_size(self):
        return np.prod(self.output_size)

    def get_input_flattened_size(self):
        return np.prod(self.output_size)*10

    def forward(self, x):
        x = x.reshape(x.shape[0],-1)
        x = self.computation(x)
        output_shape = [x.shape[0]] + self.output_size
        x = x.reshape(output_shape)
        return x


class TerpretProblem_NR(TerpretProblem):
    def __init__(self,opts):
        super(TerpretProblem_NR, self).__init__(opts)

        self.opts=opts
        self.reparameterization_net = Reparameterization_Net(size=[opts.k-1,2])

        self.latent_input_parameter = nn.Parameter(torch.rand(self.reparameterization_net.get_input_flattened_size(),dtype=torch.float), requires_grad=True)
         # = nn.Parameter(torch.rand(opts.k-1, 1,dtype=torch.float)*0.1+0.45, requires_grad=True)
        # self.parameter = nn.Parameter(torch.rand(opts.k-1, 1,dtype=torch.float)*0.01+0.495, requires_grad=True)

    def learnable_parameters(self):
        # return [self.latent_input_parameter] + [it for it in  self.reparameterization_net.computation.parameters() ]
        return [it for it in  self.reparameterization_net.computation.parameters() ]
        # return [self.latent_input_parameter]


    def forward(self, x):
        opts=self.opts
        # latent_input_parameter=self.latent_input_parameter + (torch.rand(self.latent_input_parameter.shape)-0.5)*100
        latent_input_parameter=self.latent_input_parameter
        reparameterized_parameter = self.reparameterization_net( torch.unsqueeze(latent_input_parameter,0) )[0]



        # temp = (reparameterized_parameter+(torch.rand(reparameterized_parameter.shape)-0.5)*1)/0.01
        temp = reparameterized_parameter/0.1
        self.mus = torch.cat([x, torch.nn.functional.softmax(temp,dim=1)], 0)

        # temp = torch.cat((temp, 1-temp), dim=1)
        # temp = torch.cat((1-temp, temp), dim=1)

        # self.mus = torch.cat([x, temp], 0)
        def soft_xor_p0(i):
            j = i+1
            j = j % opts.k
            return self.mus[i, 0] * self.mus[j, 0] + self.mus[i, 1] * self.mus[j, 1]

        self.ys_eq_0 = torch.stack([ soft_xor_p0(temp) for temp in range(opts.k)],0)

        # self.result = torch.log(self.ys_eq_0).sum()
        eps=1e-7
        self.result = torch.log(self.ys_eq_0+eps).sum()
        return - self.result

class Reparameterization_Net_New(nn.Module):
    def __init__(self,list_of_variable_dimension):
        super(Reparameterization_Net_New, self).__init__()
        self.list_of_variable_dimension = list_of_variable_dimension
        self.output_size = sum([ s for s in list_of_variable_dimension]) # the size of the output shape


        self.input_dimension = 10
        self.intermediate_dimension = 100

        self.modules = []
        for s in list_of_variable_dimension:
            self.output_dimension = s
            computation = nn.Sequential(
                # nn.LayerNorm(10),
                nn.Linear(self.input_dimension,self.intermediate_dimension),
                nn.ReLU(),
                nn.Linear(self.intermediate_dimension,self.intermediate_dimension),
                nn.ReLU(),
                nn.Linear(self.intermediate_dimension,self.intermediate_dimension),
                nn.ReLU(),
                nn.Linear(self.intermediate_dimension,self.intermediate_dimension),
                nn.ReLU(),
                nn.Linear(self.intermediate_dimension,self.intermediate_dimension),
                nn.ReLU(),
                nn.Linear(self.intermediate_dimension,self.intermediate_dimension),
                nn.ReLU(),
                nn.Linear(self.intermediate_dimension,self.intermediate_dimension),
                nn.ReLU(),
                nn.Linear(self.intermediate_dimension,self.intermediate_dimension),
                nn.ReLU(),
                nn.Linear(self.intermediate_dimension,self.intermediate_dimension),
                nn.ReLU(),
                nn.Linear(self.intermediate_dimension,self.intermediate_dimension),
                nn.ReLU(),
                nn.Linear(self.intermediate_dimension,self.intermediate_dimension),
                nn.ReLU(),
                nn.Linear(self.intermediate_dimension,self.intermediate_dimension),
                nn.ReLU(),
                nn.Linear(self.intermediate_dimension,self.intermediate_dimension),
                nn.ReLU(),
                nn.Linear(self.intermediate_dimension,self.intermediate_dimension),
                nn.ReLU(),
                nn.Linear(self.intermediate_dimension,self.intermediate_dimension),
                nn.ReLU(),
                nn.Linear(self.intermediate_dimension,self.output_dimension),
                # nn.Sigmoid()
            )
            self.modules.append(computation)

    def forward(self, x):
        # x = x.reshape(x.shape[0],-1)
        output = []
        for i,variable in enumerate(x):
            # print(x.shape)
            variable = self.modules[i](variable)
            output.append(variable)
        return torch.stack(output)

class TerpretProblem_NR_New(TerpretProblem):
    def __init__(self,opts):
        super(TerpretProblem_NR_New, self).__init__(opts)

        self.opts=opts
        self.reparameterization_net = Reparameterization_Net_New( [1 for _ in range(opts.k-1)] )

        self.latent_input_parameter = nn.Parameter(torch.rand(opts.k-1,10,dtype=torch.float), requires_grad=True)
         # = nn.Parameter(torch.rand(opts.k-1, 1,dtype=torch.float)*0.1+0.45, requires_grad=True)
        # self.parameter = nn.Parameter(torch.rand(opts.k-1, 1,dtype=torch.float)*0.01+0.495, requires_grad=True)

    def learnable_parameters(self):
        # return [self.latent_input_parameter] + [it for it in  self.reparameterization_net.computation.parameters() ]
        param = []
        for it in  self.reparameterization_net.modules:
            param += it.parameters()
        # param += [self.latent_input_parameter]
        return param
        # return [self.latent_input_parameter]

    def forward(self, x):
        opts=self.opts
        latent_input_parameter=self.latent_input_parameter - (torch.rand(self.latent_input_parameter.shape)-0.5)*100
        latent_input_parameter = self.latent_input_parameter

        reparameterized_parameter = self.reparameterization_net(latent_input_parameter)
        temp = reparameterized_parameter
        # print(temp.size())
        # self.mus = torch.cat([x, torch.nn.functional.softmax(temp,dim=1)], 0)

        # print("|||".join([str(1-i.tolist()[0]) for i in temp.data]))
        temp = nn.Sigmoid()(temp)
        # temp = binarize(temp)


        # temp = torch.cat((temp, 1-temp), dim=1)
        temp = torch.cat((1-temp, temp), dim=1)
        self.mus = torch.cat([x, temp], 0)

        def soft_xor_p0(i):
            j = i+1
            j = j % opts.k
            return self.mus[i, 0] * self.mus[j, 0] + self.mus[i, 1] * self.mus[j, 1]

        self.ys_eq_0 = torch.stack([ soft_xor_p0(temp) for temp in range(opts.k)],0)

        # self.result = torch.log(self.ys_eq_0).sum()
        eps=1e-7
        self.result = torch.log(self.ys_eq_0+eps).sum()
        return - self.result


def unravel_index(
    indices: torch.LongTensor,
    shape: Tuple[int, ...],
) -> torch.LongTensor:
    r"""Converts flat indices into unraveled coordinates in a target shape.

    This is a `torch` implementation of `numpy.unravel_index`.

    Args:
        indices: A tensor of (flat) indices, (*, N).
        shape: The targeted shape, (D,).

    Returns:
        The unraveled coordinates, (*, N, D).
    """

    coord = []

    for dim in reversed(shape):
        coord.append(indices % dim)
        indices = indices // dim

    coord = torch.stack(coord[::-1], dim=-1)

    return coord

class NewOp(Optimizer):
    def __init__(
        self,
        binary_params,
        adaptive_rate=0.1, # this is just placeholder
        continuous_optimizer = None
    ):
        self.total_weights = {}
        if continuous_optimizer is None:
            self.use_non_binary=False
        else:
            self.use_non_binary=True
            self._adam = continuous_optimizer

        defaults = dict(adaptive_rate=adaptive_rate)
        super(NewOp, self).__init__(
            binary_params, defaults
        )

    def step(self, closure: Optional[Callable[[], float]] = ..., ar=None):
        if self.use_non_binary is True:
            self._adam.step()

        flips = {None}

        for group in self.param_groups:
            params = group["params"]

            flips = {}

            for param_idx, p in enumerate(params):
                grad = p.grad.data

                state = self.state[p]

                ar = group["adaptive_rate"]

                if len(state) == 0:
                    state['step'] = 0
                    state["moving_average"] = torch.zeros_like(p.data, memory_format=torch.preserve_format)
                    # state["moving_average"] = torch.zeros_like(torch.rand(p.data.size()), memory_format=torch.preserve_format)

                # m is the average of gradient
                m = state['moving_average']
                m.mul_((1 - ar))
                m.add_(grad.mul(ar))
                m_=m

                # p_temp = torch.where(p>0.5, torch.ones_like(p), - torch.ones_like(p))
                #
                # mask = (m_.abs() >= t ) * (m_.sign() == p_temp.sign())
                # mask = mask.float() * 1
                # flips[param_idx] = mask
                #
                # p.data.logical_xor_(mask).type(torch.float)




                # signal= (2*p.data-1) * grad
                signal= (2*p.data-1) * m_

                temperature=0.1
                signal = signal.flatten()
                if torch.all(signal <=0):
                    continue
                signal = torch.clamp(signal,min=0)
                signal= signal/temperature
                signal = F.softmax(signal,dim=0)
                # temp=torch.argmin(signal)
                # flip_index=unravel_index(temp,grad.shape)
                # flip_index=flip_index.squeeze()
                K=5
                flip_index = torch.multinomial(signal,K,replacement=True)
                flip_index = flip_index.unique()
                # print(flip_index)


                # for i in range(p.data.shape[0]):
                #     # print( -(2*p.data[i]-1)*grad[i],p.data[i])
                #     print(signal[i],p.data[i])
                #
                # print(grad.shape)
                # print(p.data.shape)
                # print(torch.argmax(grad))
                # print(flip_index)
                for i in range(flip_index.shape[0]):
                    target_dim = flip_index[i]
                    target=p.data[target_dim,0]
                    if target == 1:
                        # p.data[flip_index[0],flip_index[1]-1]=1
                        p.data[target_dim,0]=0
                    else:
                        # p.data[flip_index[0],flip_index[1]-1]=0
                        p.data[target_dim,0]=1

                # for i in range(p.data.shape[0]):
                    # print(- p.data[i]*grad[i],p.data[i])
                # quit()

            # print('# of flip',flips)
        return

    def zero_grad(self) -> None:
        super().zero_grad()
        if self.use_non_binary is True:
            self._adam.zero_grad()