Bit operation of cuda tensor in pytorch

Viewed 120

I'd like to do some bit operation with torch.tensor in cuda, such as <<, >>, or extract each bit which represents a float number, such as 0 01101 01010101012 (3555H) for 0.333 in float16.

what I do now is as below:

def _decompose(self, value, exp_bias=None):
    '''
    decompose a single into sign, exp and mant
    '''
    if exp_bias is None:
        exp_bias = self.exp_bias
    # smallest non-zero float point
    descriminator = torch.tensor((2 ** (-exp_bias)) / 2).type_as(value)
    sign = (value > descriminator).type_as(value)
    sign -= (value < -descriminator).type_as(value)
    value = value.abs()
    exp = torch.log2(value).floor()
    mant = value / (2 ** exp)
    return sign, exp, mant

Is there a way to achieve such function? Or is there any wrong with my code? Thanks.

0 Answers
Related