Repository: pytorch
Issue: Support Dict Parameter Type for custom_op
Labels: triaged,module: custom-operators,oncall: pt2,module: pt2-dispatcher

### 🐛 Describe the bug

Is it possible to support infer_schema for custom_op which has Dict as input parameters? I think opschema can support such sig `(Tensor t, Dict(str, Any) meta) -> Tensor`. Also, can such inputs be mutated?

```python
import torch
from typing import Dict, Any

@torch.library.custom_op("host_code::collect_max", mutates_args=(), device_types="cpu")
def fn(t: torch.Tensor, meta: Dict[str, Any]) -> torch.Tensor:
    meta["max"] = t.max().item()
    return t.clone()

@torch.library.register_fake("host_code::collect_max")
def fn_fake(t: torch.Tensor, meta: Dict[str, Any]) -> torch.Tensor:
    return t


t = torch.randn((3,3))
meta = {}
fn(t, meta)
print(meta)
```

### Error logs

```
Traceback (most recent call last):
  File "/Users/chenxiny/workspace/dynamo_case/custom_op.py", line 5, in <module>
    def fn(t: torch.Tensor, meta: Dict[str, Any]):
  File "/Users/chenxiny/miniforge3/envs/torch-metal/lib/python3.10/site-packages/torch/_library/custom_ops.py", line 121, in inner
    schema_str = torch.library.infer_schema(fn, mutates_args=mutates_args)
  File "/Users/chenxiny/miniforge3/envs/torch-metal/lib/python3.10/site-packages/torch/_library/infer_schema.py", line 106, in infer_schema
    error_fn(
  File "/Users/chenxiny/miniforge3/envs/torch-metal/lib/python3.10/site-packages/torch/_library/infer_schema.py", line 58, in error_fn
    raise ValueError(
ValueError: infer_schema(func): Parameter meta has unsupported type typing.Dict[str, typing.Any]. The valid types are: dict_keys([<class 'torch.Tensor'>, typing.Optional[torch.Tensor], typing.Sequence[torch.Tensor], typing.List[torch.Tensor], typing.Sequence[typing.Optional[torch.Tensor]], typing.List[typing.Optional[torch.Tensor]], <class 'int'>, typing.Optional[int], typing.Sequence[int], typing.List[int], typing.Optional[typing.Sequence[int]], typing.Optional[typing.List[int]], <class 'float'>, typing.Optional[float], typing.Sequence[float], typing.List[float], typing.Optional[typing.Sequence[float]], typing.Optional[typing.List[float]], <class 'bool'>, typing.Optional[bool], typing.Sequence[bool], typing.List[bool], typing.Optional[typing.Sequence[bool]], typing.Optional[typing.List[bool]], <class 'str'>, typing.Optional[str], typing.Union[int, float, bool], typing.Union[int, float, bool, NoneType], typing.Sequence[typing.Union[int, float, bool]], typing.List[typing.Union[int, float, bool]], <class 'torch.dtype'>, typing.Optional[torch.dtype], <class 'torch.device'>, typing.Optional[torch.device]]). Got func with signature (t: torch.Tensor, meta: Dict[str, Any]))
```

### Versions

```
PyTorch version: 2.6.0.dev20241126
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 15.1.1 (arm64)
GCC version: Could not collect
Clang version: 15.0.0 (clang-1500.3.9.4)
CMake version: Could not collect
Libc version: N/A

Python version: 3.10.15 | packaged by conda-forge | (main, Oct 16 2024, 01:24:20) [Clang 17.0.6 ] (64-bit runtime)
Python platform: macOS-15.1.1-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Apple M1 Pro

Versions of relevant libraries:
[pip3] numpy==2.1.2
[pip3] torch==2.6.0.dev20241126
[pip3] torchaudio==2.5.0.dev20241126
[pip3] torchvision==0.20.0.dev20241126
[conda] numpy                     2.1.2                    pypi_0    pypi
[conda] torch                     2.6.0.dev20241126          pypi_0    pypi
[conda] torchaudio                2.5.0.dev20241126          pypi_0    pypi
[conda] torchvision               0.20.0.dev20241126          pypi_0    pypi
```

cc @chauhang @penguinwu @zou3519 @bdhirsh @yf225