def argument(
a: Union[Argument, TensorOptionsArguments,
ThisArgument]) -> CppArgument:
if isinstance(a, Argument):
return CppArgument(
type=argument_type(a),
name=a.name,
default=default_expr(a.default, a.type)
if a.default is not None else None,
argument=a,
)
elif isinstance(a, ThisArgument):
return CppArgument(
type=argument_type(a.argument),
name=
"const_cast<Tensor&>(*this)", # this is an abuse but it's convenient
default=None,
argument=a,
)
elif isinstance(a, TensorOptionsArguments):
default = None
if all(x.default == "None" for x in a.all()):
default = '{}'
elif a.dtype.default == "long":
default = 'at::kLong' # TODO: this is wrong
return CppArgument(
type='const TensorOptions &',
name='options',
default=default,
argument=a,
)
else:
assert_never(a)
def exprs(args: Sequence[DispatcherArgument]) -> Sequence[DispatcherExpr]:
return cpparguments_exprs([
CppArgument(type=a.type,
name=a.name,
default=None,
argument=a.argument) for a in args
])
def exprs(args: Sequence[DispatcherArgument]) -> Sequence[DispatcherExpr]:
if local.use_c10_dispatcher() is UseC10Dispatcher.full:
process_tensoroptions = ProcessTensoroptions.SCATTER
else:
process_tensoroptions = ProcessTensoroptions.PASS_THROUGH
return cpparguments_exprs([
CppArgument(
type=a.type, name=a.name, default=None, argument=a.argument)
for a in args
],
process_tensoroptions=process_tensoroptions)