mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/49042 I want the names positional and kwarg_only to give the unflat representation (e.g., preserving TensorOptionsArguments in the returned Union). So I regret my original naming choice when I moved grouping to model. This renames them to have flat_ prefix and also adds a flat_non_out argument for cases where you just want to look at non-out arguments. Signed-off-by: Edward Z. Yang <ezyang@fb.com> Test Plan: Imported from OSS Reviewed By: H-Huang Differential Revision: D25455884 Pulled By: ezyang fbshipit-source-id: f923f8881267a3e3e8e9521519412f7cc25034fc
32 lines
1.0 KiB
Python
32 lines
1.0 KiB
Python
from tools.codegen.model import *
|
|
from tools.codegen.api.types import MetaArgument
|
|
|
|
import tools.codegen.api.dispatcher as dispatcher
|
|
|
|
from typing import Sequence
|
|
|
|
# Follows dispatcher calling convention, but:
|
|
# - Mutable arguments not allowed. Meta functions are always
|
|
# written in functional form. Look at FunctionSchema.signature()
|
|
# - No tensor returns; instead we return a TensorMeta describing
|
|
# the tensor in question
|
|
|
|
def name(g: StructuredNativeFunctions) -> str:
|
|
# use the overload name from the functional version
|
|
return str(g.functional.func.name).replace('.', '_')
|
|
|
|
def argument_type(a: Argument) -> str:
|
|
assert not a.is_write
|
|
return dispatcher.argumenttype_type(a.type, mutable=False)
|
|
|
|
def argument(a: Argument) -> MetaArgument:
|
|
return MetaArgument(
|
|
type=argument_type(a),
|
|
name=a.name,
|
|
argument=a,
|
|
)
|
|
|
|
def arguments(func: FunctionSchema) -> Sequence[MetaArgument]:
|
|
assert not func.arguments.out
|
|
return list(map(argument, func.arguments.flat_non_out))
|