code: def test_unbind(dim: int, contiguous: bool):
      x = torch.randn([10, 20, 4, 10, 3])
      x2 = x.clone()
  
      if not contiguous:
          perm = list(range(x.ndim))
          random.Random(dim).shuffle(perm)
          # Let's hope we didn't pick identity
          x = x.permute(perm)
          x2 = x2.permute(perm)
      assert contiguous == x.is_contiguous()
      x.requires_grad_(True)
      x2.requires_grad_(True)
  
      # FW
      tensors = xformers.ops.unbind(x, dim)
      tensors2 = torch.unbind(x2, dim)
      assert len(tensors) == len(tensors2)
      for t1, t2 in zip(tensors, tensors2):
          assert torch.allclose(t1, t2)
  
      # BW
      grads = torch.unbind(torch.randn(x.shape), dim)
      zero = torch.zeros_like(tensors[0])
      loss1 = sum(((g * t) for (g, t) in zip(grads, tensors)), zero)
      loss2 = sum(((g * t) for (g, t) in zip(grads, tensors2)), zero)
      assert torch.allclose(loss1, loss2)
      g = torch.randn_like(loss1)
      loss1.backward(g)
      loss2.backward(g)
      assert x.grad is not None
      assert x2.grad is not None
      assert torch.allclose(x.grad, x2.grad)
  
imports: import random
  import pytest
  import torch
  import xformers.ops
  from xformers.ops.common import _get_storage_base
label_desc: Use image node_id 1 for calling a global function with example usage: test_unbind(dim, contiguous) without return types