Repository: pytorch
Issue: RFC: [pytree] node registration namespaces
Labels: triaged,needs design,module: pytree

### 🚀 The feature, motivation and pitch

We should define (and freeze) pytree node registration namespaces. External packages must define their own namespace.

torch namespaces:
- **"torch"** - default datatypes (list, dict, OrderedDict, namedtuple, tuple)
    - `hygeine_level="ownership"` - any registrations that satisfy ownership.
- **"torch.vmap.Size"** - custom `torch.Size` behaviour - unflatten as tuple _if any elem isinstance(NoneType)_
- **"torch.fx.proxy_tensor.Size"** -  custom `torch.Size` behaviour - unflatten as tuple _if any elem isinstance(Proxy)_
- **"torch.export"** - registering custom dataclasses. 
    - `hygeine_level="external,ownership"`; external packages requiring the dataclass pytree in their non-torch.export code must run `@using_pytree_namespace("torch.export")`
    
Open to debate:
- **"torch.fx.immutable_collections"** - this one is debatable. Since there is sufficient hygeine - class definition module is registration module, "torch" namespace might be fine (allow ambient registration if you own the class). Provisionally, we will allow under `torch` namespace's due to `hygeine_level="ownership"`

### Default namespace
Default namespace is `torch`

### Freeze API / Hygeine Level API
There will be 3 levels of hygeine. The hygeine level will be enforced outside of the module in which the namespace is defined (implict ownership of namespace). Within the same module, it is possible to register as many classes as you like.

1. `hygeine_level="freeze"`: totally frozen, cannot add any new classes no matter what
2. `hygeine_level="ownership"`: can add new classes, but only if the registration is performed in the same python module in which the class is defined
3. `hygeine_level="external"`: can add new classes, but only if the registration is not performed within `torch.*` package

We will never allow adding new classes within `torch.*` without restriction.
`torch` namespace will be created with `create_namespace("torch", hygeine_level="ownership")`.
`"ownership,external"` is a valid hygeine level - either the class is owned, or external.

### Namespace inheritance

Context manager style is no longer proposed. See discussion below
<details>

Unclear if `torch.fx.proxy_tensor` namespace should inherit from `torch.fx`. 

We say no - ideally
1. inheritence of namespace is only from default (and not from any notion of string path), and can be opted out. `pytree.create_namespace(namespace="?", use_default_registrations=True)`. 
2. Else, inner inherits from outer inherits from ambient unless overriden. If we want inner-overrides-outer, then we must also allow inner to unset a given class as a non-leaf node. `pytree.create_namespace(namespace="?", inherit_ambient_registrations=True)`.

</details>

### Combining namespaces

Context manager style is no longer proposed. See discussion below
<details>

It is unclear if there can be nested/parallel namespaces. We can define nested/parallel namespace, but either:
1. there needs to be validation that `namespace1.registered_nodes.isdisjoint(namespace2.registered_nodes)`
3. inner overrides outer (this is consistent with how `@using_pytree_namespace` can override ambient)

```python
@using_pytree_namespace(["torch.export", "torch.vmap.Size"])

@using_pytree_namespace("torch.export")
def fn(x):
    @using_pytree_namespace("torch.vmap.Size")
    def g(x):
        ..
```

</details>

### Locally Enabling Namespaces

```python
def fn():
     _, _ = pytree.tree_flatten(.., namespace="torch.vmap.Size")
```


Context manager style is no longer proposed. See discussion below
<details>

```python
@using_pytree_namespace("torch.vmap.Size")
def fn():
     _, _ = pytree.tree_flatten(..)

```

```python
def fn():
    with using_pytree_namespace("torch.vmap.Size"):
        _, _ = pytree.tree_flatten(..)
```

</details>

### Handling Errors
Errors should be namespace aware. See https://github.com/pytorch/pytorch/issues/112342#issuecomment-1784227795

### Inspecting Namespaces

```python
print(pytree.show_namespace("torch.vmap.Size"))
# { 'registrations': 
#     [
#         {"cls": torch.Size, "flatten_fn": <lambda @ x13f713>, "unflatten_fn": <lambda @ xha3f713> }
#         {"cls": list, ...}
#         {"cls": dict, ...}
#     ],
#    'created_at': 'torch/_functorch/vmap.py',
#    'inherit_ambient_registrations': True,
#    'hygeine_level': 'freeze'
# }
```

cc: @XuehaiPan @peterbell10 @zou

### Alternatives

NIL

### Additional context

See discussion: 
1. https://github.com/pytorch/pytorch/issues/112342
4. https://github.com/pytorch/pytorch/pull/112111

cc @zou3519