Bitwise Operations on Cuda Float Tensor

mmackay: <p>I would like to access the bit representation of a float tensor on a GPU and perform manipulations such as shifting, anding, etc. on it. I am wondering the best way to do this and if it’s even possible in PyTorch.  Thanks for the help!</p>

xtipacko: <p>I also need this feature!</p>

smth: <p>We only support bitwise operations on Integer Tensor types at the moment.<br>
The supported operations are listed here:</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/master/torch/csrc/generic/methods/Tensor.cwrap#L912-L1120" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/master/torch/csrc/generic/methods/Tensor.cwrap#L912-L1120" target="_blank" rel="nofollow noopener">pytorch/pytorch/blob/master/torch/csrc/generic/methods/Tensor.cwrap#L912-L1120</a></h4>
<pre class="onebox"><code class="lang-cwrap"><ol class="start lines" start="912" style="counter-reset: li-counter 911 ;">
<li>return: void*</li>
<li>cpu_half: True</li>
<li>cname: data</li>
<li>arguments:</li>
<li>  - THTensor* self</li>
<li>]]</li>
<li>
</li>
<li>[[</li>
<li>name: equal</li>
<li>variants:</li>
<li>  - method</li>
<li>  - function</li>
<li>return: bool</li>
<li>arguments:</li>
<li>  - THTensor* self</li>
<li>  - THTensor* other</li>
<li>]]</li>
<li>
</li>
<li>[[</li>
<li>python_name: copy_</li>
</ol></code></pre>

  This file has been truncated. <a href="https://github.com/pytorch/pytorch/blob/master/torch/csrc/generic/methods/Tensor.cwrap#L912-L1120" target="_blank" rel="nofollow noopener">show original</a>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<p>You can write your own C/C++ function that takes in the CUDA float* pointer and you can manipulate it yourself. Here’s an example of writing a CUDA extension: <a href="https://github.com/longcw/yolo2-pytorch/tree/master/layers/reorg">https://github.com/longcw/yolo2-pytorch/tree/master/layers/reorg</a></p>
<p>Alternatively, you can use the <code>cupy</code> package, and write your whole CUDA kernel / manipulation in python scripts, for example: <a href="https://github.com/szagoruyko/pyinn/blob/master/pyinn/conv2d_depthwise.py">https://github.com/szagoruyko/pyinn/blob/master/pyinn/conv2d_depthwise.py</a></p>
<p>What exactly is your use-case?</p>

Geremia: <p>Can’t you <a href="https://pytorch.org/docs/stable/autograd.html#function" rel="noopener nofollow ugc">subclass <code>torch.autograd.Function</code></a> and mark it as <a href="https://pytorch.org/docs/stable/generated/torch.autograd.function.FunctionCtx.mark_non_differentiable.html" rel="noopener nofollow ugc"><code>torch.autograd.function.FunctionCtx.mark_non_differentiable</code></a>?</p>

Geremia: <p>Something like this?</p>
<pre data-code-wrap="python"><code class="lang-python">class BitShiftFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, shift):
        ctx.x = x
        ctx.shift = shift.item()
        x_int = x.view(torch.int)
        if ctx.shift &gt; 0:
            shifted = x_int &lt;&lt; ctx.shift
        else:
            shifted = x_int &gt;&gt; -ctx.shift
        return shifted.view(torch.float32)

    @staticmethod
    def backward(ctx, grad_output):
        x = ctx.x
        shift = ctx.shift
        ctx.mark_non_differentiable(shift)  # Mark shift as non-differentiable
        return x, None

# Usage example
shift = 2
x = torch.randn(3, 3, requires_grad=True)
result = BitShiftFunction.apply(x, torch.tensor(shift))
result.backward(torch.ones_like(result))

print(f"x:\n{x}")
print(f"Shifted x:\n{result}")
print(f"Gradients of x:\n{x.grad}")
</code></pre>
<pre data-code-wrap="shell"><code class="lang-shell">x:
tensor([[-2.5461, -0.0995,  0.4327],
        [-0.5932, -0.4999, -0.7176],
        [-0.7571,  0.7862,  1.7181]], requires_grad=True)
Shifted x:
tensor([[ 1.2839e-38, -3.5585e+33, -1.2778e+36],
        [-4.6406e+36, -2.6573e+36, -9.2560e+36],
        [-1.1241e+37, -1.3716e+37, -3.1861e+38]],
       grad_fn=&lt;BitShiftFunctionBackward&gt;)
Gradients of x:
tensor([[-2.5461, -0.0995,  0.4327],
        [-0.5932, -0.4999, -0.7176],
        [-0.7571,  0.7862,  1.7181]])
</code></pre>
<p><sup><img src="https://discuss.pytorch.org/images/emoji/apple/tophat.png?v=12" title=":tophat:" class="emoji" alt=":tophat:" loading="lazy" width="20" height="20">-tip: based off <a href="https://hf.co/chat/r/NMQsnI1?leafId=3f5fba3b-9fff-4527-ad69-3bfe91f2cc37" rel="noopener nofollow ugc">this Deepseek answer</a></sup></p>

ptrblck: <p>Cross-posting my response to this approach: <a href="https://discuss.pytorch.org/t/bitwise-operation-on-float-tensor/170863/11" class="inline-onebox">Bitwise Operation on Float Tensor - #11 by ptrblck</a></p>