Repository: pytorch
Issue: Scalar/Tensor arg type for op schemas
Labels: triaged,module: codegen

Defining an op that supports passing either a Scalar or Tensor for one or more args is currently laborious - there's no way to type an arg as taking both (i.e., a union), instead one must write out each variant as a distinct overload. When an op has more than such arg, the situation gets untenably combinatorial, esp. when combined with other such multipliers (in/func/out variants in particular). 

We want to make a new type available for op schemas that denotes "either a Tensor, or a Scalar that will be promoted to a Tensor" - i.e., C++ kernels can be defined to take Tensor at these positions; we'll take care of wrapping the scalar in a Tensor in either of the three callsite-to-kernel pipelines (eager, JIT, C++ API). 

(Note: for ops that want to preserve the Tensor/Scalar distinction and write different kernels for each variant, the option of writing separate overloads in `native_functions.yaml` will of course still be available.)

---

Aside:

It turns out that support for autoboxing (as it were) Scalars into Tensors is [already available for eager op bindings](https://github.com/pytorch/pytorch/blob/master/torch/csrc/utils/python_arg_parser.cpp#L1059-L1092), but is limited by fiat to a [specific set of ops](https://github.com/pytorch/pytorch/blob/master/torch/csrc/utils/python_arg_parser.cpp#L68-L86).

To muddy the waters further, many of these blessed ops _also_ have separate Scalar and Tensor overloads defined in native_functions. The Scalar variants are used by [macro-generated C++ operator overloads](https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/TensorOperators.h#L82-L96) that make the full set of Scalar/Tensor combinations available in the C++ API. (@ezyang also mentions seeing JIT test errors when Scalar variants are removed from `native_functions.yaml`, implying that they're also relied on there, but I'm not sure how.)

---

The goal is to 
* provide support for any operator to declare one or more `ScalarOrTensor` (placeholder name) args in native_functions
  * we also need `ScalarOrTensor?` for e.g. [np.diff](https://github.com/pytorch/pytorch/pull/50569); I'm sure uses cases for `ScalarOrTensor[]` etc. will show up soon enough
* trigger support purely from this specification (i.e., no hidden whitelists)
* provide uniform support across eager, JIT, and C++ APIs

To do this we'll need to do the following.
* In codegen:
  * add a new base type `ScalarOrTensor` (or whatever). 
  * generated eager bindings probably just need to pass the type through in generated schema strings; the necessary autoboxing behavior will be triggered when a passed Scalar is matched against a ScalarOrTensor-typed param. 
  * similarly, the JIT infra should just be triggered by the type name appearing in schema
  * for generated C++ API artifacts, we'll need to generate overloads to support the different combinations. We can generate full implementations for the Scalar variants that (like the eager and JIT call resolution machinery) use `wrapped_scalar_tensor` to box and then call the Tensor overloads.
* In eager binding infra (python_arg_parser): add logic to trigger the scalar-to-tensor wrapping from the new type.
* In JIT: add logic to properly compile callsites to ops with params of this type: passthrough for Tensor-typed args, wrapping for Scalar-typed args, conditional wrapping for dynamically typed (or union!) args. (cc @eellison, is there any broadly similar implicit-promotion type stuff being done that we could pattern this off of?)
* In C++: port or remove the macros in `TensorOperators.h` that currently generate the full set of covering overloads (implemented by calling the subset of handwritten Scalarized overloads from native_functions on swizzled args), since now we'll be doing this from codegen

After support is built, we can
* update the hand-overloaded ops already in `native_functions.yaml` to use the new type
* remove the now-unnecessary whitelist and the logic that uses it from python_arg_parser.cpp
* update any other ops in `native_functions.yaml` that are currently forgoing Scalar support in anticipation of this behavior becoming available (e.g. [np.diff](https://github.com/pytorch/pytorch/pull/50569)