[batching] pixel_{un}shuffle

[kshitij12345]

Reference: #240

[kshitij12345]

Failure looks unrelated

test_make_fx_exhaustive_linalg_lstsq_grad_oriented_cpu_float32

Traceback (most recent call last):
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/testing/_internal/common_device_type.py", line 376, in instantiated_test
    result = test(self, **param_kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/testing/_internal/common_device_type.py", line 777, in dep_fn
    return fn(slf, *args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/testing/_internal/common_device_type.py", line 777, in dep_fn
    return fn(slf, *args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/testing/_internal/common_device_type.py", line 742, in test_wrapper
    return test(*args, **kwargs)
  File "test/test_pythonkey.py", line 249, in test_make_fx_exhaustive
    self.assertEqual(new_out, old_out)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/testing/_internal/common_utils.py", line 1985, in assertEqual
    super().assertTrue(result, msg=self._get_assert_msg(msg, debug_msg=debug_msg))
AssertionError: False is not true : Tensors failed to compare as equal!With rtol=1.3e-06 and atol=1e-05, found 3 element(s) (out of 81) whose difference(s) exceeded the margin of error (including 0 nan comparisons). The greatest difference was 2.5272369384765625e-05 (5.405212879180908 vs. 5.405187606811523), which occurred at index (0, 0, 0, 1).

[Chillee]

LGTM!