Repository: pytorch
Issue: Adding a new kwarg to a torch.nn.functional function breaks FC for JIT 
Labels: oncall: jit,module: nn,triaged,has workaround,better-engineering

## 🐛 Bug

Even though #54613 has been addressed, the problem of JIT FC breakage when adding a new default-valued kwarg still occurs for functions in `torch.nn.functional`.

To elucidate this, here's an example real-world FC-breaking issue I ran into when attempting to add a new kwarg `label_smoothing=0.0` to `torch.nn.functional.cross_entropy`. Here's the Python-side code; the functional wrapper `torch.nn.functional.cross_entropy` calls the underlying op `cross_entropy_loss`:

```python
def cross_entropy(
    input: Tensor,
    target: Tensor,
    weight: Optional[Tensor] = None,
    size_average: Optional[bool] = None,
    ignore_index: int = -100,
    reduce: Optional[bool] = None,
    reduction: str = "mean",
    label_smoothing: float = 0.0,
) -> Tensor:
    # __torch_function__ stuff omitted for brevity
    ...
    if size_average is not None or reduce is not None:
        reduction = _Reduction.legacy_get_string(size_average, reduce)
    return torch._C._nn.cross_entropy_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index, label_smoothing)
```

The serialized form of the JIT-scripted functional wrapper in `code/__torch__/torch/nn/functional.py` looks like this; note that the new kwarg is unconditionally present in the call to the underlying op (bad news for FC).

```python
def cross_entropy(input: Tensor,
    target: Tensor,
    weight: Optional[Tensor]=None,
    size_average: Optional[bool]=None,
    ignore_index: int=-100,
    reduce: Optional[bool]=None,
    reduction: str="mean",
    label_smoothing: float=0.) -> Tensor:
  _0 = __torch__.torch.nn._reduction.legacy_get_string
  _1 = __torch__.torch.nn._reduction.get_enum
  if torch.__isnot__(size_average, None):
    _2, size_average0 = True, unchecked_cast(bool, size_average)
  else:
    _2, size_average0 = torch.__isnot__(reduce, None), size_average
  if _2:
    reduction0 = _0(size_average0, reduce, True, )
  else:
    reduction0 = reduction
  _3 = torch.cross_entropy_loss(input, target, weight, _1(reduction0, ), ignore_index, label_smoothing)
  return _3
```

With the use of the wrapper, a non-default value for `label_smoothing` could be passed to the underlying op, so this is unfortunate but expected.

The general pattern for modules in `torch.nn` is to always route through such functional wrappers, so the FC problem realistically remains for any new default-valued kwarg added within `torch.nn` (and potentially other places that employ op wrappers).

### Possible Solutions
* Remove the `torch.nn.functional` wrappers (e.g. have `F.foo` *be* the underlying op function). This would be a decent amount of work, and maybe impossible for some functions (needs additional exploration). Arguably, these wrappers just add bloat to the codebase and not much else in a lot of cases, so this is probably a good ideal to strive for anyway.
* Change the JIT to be aware of these wrappers and handle them specially? I don't think this is a realistic solution; there are lots of cons here coming from the additional coupling..
* ???

### Workaround
For now, the following workaround makes it possible add a new, default-value kwarg to a `torch.nn` module / function:
1. Add the new default-valued kwarg to the op signature only
2. Wait the 2-week FC window
3. Update the functional and module wrappers in `torch.nn` to support the new kwarg

cc @albanD @mruberry @jbschlosser @walterddr