[BUG] functorch.dim breaks split of big tensors

[vmoens]

Describe the bug

For some reason (probably unrealted to tensordict), functorch dims fails with big tensors.

import torch
from tensordict import TensorDict

# build splits
total = 0
l = []
while total < 6400:
    l.append(torch.randint(2, 10, (1,)).item())
    total += l[-1]

# build td
tensordict = TensorDict({"x": torch.randn(6401, 1), "y": torch.randn(6401, 1)}, [])

# split
for k, x in tensordict.items():
    x.split(l, 0)

This should result in an error like this:

libc++abi: terminating with uncaught exception of type c10::Error: allocated_ <= ARENA_MAX_SIZE INTERNAL ASSERT FAILED at "/stuff/pytorch/pytorch/pytorch/functorch/csrc/dim/arena.h":227, please report a bug to PyTorch. 

[tcbegley]

I don't understand this yet, but can confirm that this appears to be unrelated to tensordict.

import torch

t = torch.rand(6401, 1)

# build splits
total = 0
l = []
while total < 6400:
    l.append(torch.randint(2, 10, (1,)).item())
    total += l[-1]

# could get RuntimeError if `sum(l) != 6401`
l[-1] = 6401 - sum(l[:-1])

# this is fine
_ = t.split(l, 0)

But now, if we import functorch.dim it will fail

import functorch.dim

_ = t.split(l, 0)

So something is happening at import with functorch that results in this bad behaviour. I'll try to pin it down a bit more but seems clear it's an issue upstream.

[vmoens]

see this bug for ref