Repository: pytorch
Issue: Functorch pytrees with custom iterables
Labels: triaged,enhancement,module: pytree,module: functorch

### 🚀 The feature, motivation and pitch

`torch.func` / `functorch` does support some sort of pytrees similar to JAX. However, it only seems to accept Lists and Dicts as iterable inputs.

Example:

```
import torch.func as FT
import torch


class DotDict(dict):
    """dot.notation access to dictionary attributes"""

    __getattr__ = dict.get
    __setattr__ = dict.__setitem__
    __delattr__ = dict.__delitem__


a = DotDict({
    "A": torch.randn(2,2),
    "inner": {
        "B": torch.randn(2,2),
        "C": torch.randn(2,2)
        }
    })

b = [
    torch.randn(2,2),
    torch.randn(2,2),
]

def f(x, y):
    return (x.A + x.inner["B"] @ x.inner["C"]).sum() + y[0].sum() + y[1].sum()


grad, val = FT.grad_and_value(f, argnums=(0,1))(a, b)

print(grad)
```

This doesn't work.
```
Thing passed to transform API must be Tensor, got <class '__main__.DotDict'>
  File "/home/cornelius/anaconda3/envs/general/lib/python3.10/site-packages/torch/_functorch/eager_transforms.py", line 55, in create_differentiable
    raise ValueError(f'Thing passed to transform API must be Tensor, '
  File "/home/cornelius/anaconda3/envs/general/lib/python3.10/site-packages/torch/utils/_pytree.py", line 196, in <listcomp>
    return tree_unflatten([fn(i) for i in flat_args], spec)
  File "/home/cornelius/anaconda3/envs/general/lib/python3.10/site-packages/torch/utils/_pytree.py", line 196, in tree_map
    return tree_unflatten([fn(i) for i in flat_args], spec)
  File "/home/cornelius/anaconda3/envs/general/lib/python3.10/site-packages/torch/_functorch/eager_transforms.py", line 57, in _create_differentiable
    return tree_map(create_differentiable, inps)
  File "/home/cornelius/anaconda3/envs/general/lib/python3.10/site-packages/torch/_functorch/pytree_hacks.py", line 12, in <listcomp>
    [fn_(arg) for arg in flat_args]
  File "/home/cornelius/anaconda3/envs/general/lib/python3.10/site-packages/torch/_functorch/pytree_hacks.py", line 12, in tree_map_
    [fn_(arg) for arg in flat_args]
  File "/home/cornelius/anaconda3/envs/general/lib/python3.10/site-packages/torch/_functorch/eager_transforms.py", line 1243, in wrapper
    tree_map_(partial(_create_differentiable, level=level), diff_args)
  File "/home/cornelius/anaconda3/envs/general/lib/python3.10/site-packages/torch/_functorch/vmap.py", line 39, in fn
    return f(*args, **kwargs)
  File "/home/cornelius/Projects/CalibrationBound/test_functorch_pytrees.py", line 30, in <module>
    grad, val = FT.grad_and_value(f, argnums=(0,1))(a, b)
  File "/home/cornelius/anaconda3/envs/general/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "/home/cornelius/anaconda3/envs/general/lib/python3.10/runpy.py", line 196, in _run_module_as_main (Current frame)
    return _run_code(code, main_globals, None,
ValueError: Thing passed to transform API must be Tensor, got <class '__main__.DotDict'>
```

However, if `a` is a `Dict` it works fine.

It would be nice, if the definition of pytrees was extended to duck-typed iterables. I do not see why the example above shouldnt work.

### Alternatives

_No response_

### Additional context

_No response_

cc @zou3519 @Chillee @samdow @soumith @kshitij12345 @janeyx99