ShapeEnv

class torch.fx.experimental.symbolic_shapes.ShapeEnv(*, should_record_events=None, tracked_fakes=None, **kwargs)

add_var_to_val(expr, val)

向符号环境中添加一个新的符号。

bind_symbols(placeholders, args)

给定符号化大小的占位符（fake tensors）和具有真实大小的具体参数（regular tensors）的成对列表，返回一个将每个符号映射到其真实值的字典。例如，如果一个占位符的大小为 (s0, s1)，将其绑定到 (2, 4) 将得到 {s0: 2, s1: 4}。这不保证绑定 ShapeEnv 中的所有符号；如果一个符号不出现在任何占位符中，我们将无法绑定它，而已经有替代项的符号也不会被绑定。

这与 evaluate_guards 有些重复，但又足够不同，因此创建另一个副本似乎是最清晰的做法。这假设 guards 已经检查过，但如果开销很小，我们会检查是否存在意外情况。

返回类型

dict[sympy.Symbol, int]

bound_sympy(expr, size_oblivious=False)

给定一个 sympy 表达式，计算它可以取值的 ValueRanges 边界。

返回类型

ValueRanges

check_equal(other)

比较另一个 ShapeEnv 是否等价。

cleanup()

断开循环引用。

这会销毁堆栈。如果你真的想保留它们，我们只需要一种方法来断开代码对象的引用。

create_symbol(val, source, dynamic_dim=DimDynamic.DUCK, constraint_dim=None, positive=True, do_not_specialize_zero_one=False, symbolic_context=None)

创建由此 ShapeEnv 跟踪的新符号。

返回类型

Expr

create_symbolic_sizes_strides_storage_offset(ex, source, *, symbolic_context=None)

为给定的张量返回符号化的尺寸和步幅列表。我们尽力用尺寸来表示步幅，以免引入新的符号变量。

返回类型

tuple[tuple[IntLikeType, …], tuple[IntLikeType, …], IntLikeType]

create_symboolnode(sym)

从 sympy 布尔表达式创建 SymBool 对象。

返回类型

SymBool

create_symfloatnode(sym, *, hint, source=None)

从符号表达式创建 SymFloat 值。

返回类型

FloatLikeType

create_symintnode(sym, *, hint, source=None)

从符号表达式创建 SymInt 值。

如果您知道要创建的 SymInt 的当前提示值，请将其传递给 hint。否则，请传递 None，我们将尽最大努力进行猜测。

返回类型

IntLikeType

create_unbacked_symbool()

创建没有提示值的符号布尔值。

返回类型

SymBool

create_unbacked_symfloat()

创建没有提示值的符号浮点数。

返回类型

SymFloat

create_unbacked_symint(source=None)

创建没有提示值的符号整数。

返回类型

SymInt

create_unspecified_symbol(val, source, dynamic_dim=DimDynamic.DUCK, constraint_dim=None, symbolic_context=None)

创建具有不确定值的符号。

与标准符号相比，我们不假定值为正，也不针对零或一值进行专门化。

返回类型

Expr

create_unspecified_symint_and_symbol(value, source, dynamic_dim)

创建包装新不确定符号的 SymInt。

返回类型

IntLikeType

deserialize_symexpr(code)

供 compile_fx 反序列化 symexpr 使用。

返回类型

Union[SymInt, SymFloat, SymBool]

evaluate_expr(orig_expr, hint=None, fx_node=None, size_oblivious=False, fallback_value=None, *, forcing_spec=False)

给定一个表达式，对其进行求值，并在必要时添加 guard。当 fallback_value 不为 None 时，函数将返回 fallback_value 而不是因数据依赖而失败。

返回类型

Basic

evaluate_guards_expression(code, args)

预期与 produce_guards_expression() 一起使用。为给定的具体参数求值由 produce_guards_expression() 生成的表达式。

返回类型

布尔值

evaluate_guards_for_args(placeholders, args, *, ignore_static=True)

为图的占位符值生成 guards 并使用 args 对 guards 进行求值。

返回类型

布尔值

evaluate_sym_node(sym_node, size_oblivious=False, fallback_value=None)

给定一个 SymNode，对其 expr 进行求值，并在必要时添加 guards。

返回类型

Basic

evaluate_symexpr(code)

供 compile_fx 求值 symexpr 使用。

返回类型

Union[int, float, bool]

format_guards(verbose=False)

根据可选的 traceback 信息格式化此 shape env 的 guard 表达式（如果 verbose 为 True）。

返回类型

str

freeze()

冻结此 ShapeEnv 以停止累积 guards。

已冻结的 ShapeEnv 将忽略在其上生成的任何后续 guards，仅发出警告，这可能导致准确性问题。

freeze_runtime_asserts()

冻结此 ShapeEnv 以停止添加延迟的运行时断言。

如果您尝试安装新的运行时断言，当它被冻结时，我们将报错。这可能表示一个 lowering 违规，或者可能是我们静态已知为 True 但以一种无法清楚证明的方式再次检查它的情况。

get_axioms(symbols=None, compute_hint=False)

给定表达式中的符号，它会返回所有与这些符号相关的运行时断言，并与所有 guards 连接。如果 symbols 为 None，则返回所有运行时断言（以及所有 guards）。

返回类型

tuple[‘sympy.logic.boolalg.Boolean’, …]

get_implications(e)

给定一个表达式，它会返回该表达式的谓词列表。

返回类型

tuple[tuple[‘sympy.logic.boolalg.Boolean’, sympy.logic.boolalg.BooleanAtom], …]

get_nontrivial_guards()

返回一组不被静态已知的 guard 表达式（即非平凡的）。

返回类型

list[‘sympy.logic.boolalg.Boolean’]

get_pruned_guards(symints)

获取 guards 列表，但进行了修剪，使其仅提供引用传入的 symints 的 guards。

返回类型

list[torch._guards.ShapeGuard]

guard_or_defer_runtime_assert(orig_expr, msg, fx_node=None)

在可以的情况下，添加一个 guard 使 orig_expr 为 True，或者回退到添加一个在运行时检查的 assert。

参数

• orig_expr (sympy.Expr) – 断言为真的布尔表达式

• msg (str) – 在断言失败时显示的错误消息

• fx_node (Optional, torch.fx.Node) – `self.graph` 中与表达式对应的节点，如果适用。

返回类型

布尔值

ignore_fresh_unbacked_symbols()

指示新分配的未备份 SymInts 将被丢弃。

返回类型

Iterator[None]

is_unbacked_symint(symbol)

检查 sympy 符号是否符合未备份符号的命名约定。

返回类型

布尔值

patch_source_specialization(source, check_fn)

暂时向 ShapeEnv 添加符号级别的公理。这在您希望“分叉”并拥有并行 ShapeEnvs 时很有用。例如，我们在进行多图编译时使用它，以便支持具有不同专业化级别的各种图。

此上下文管理器允许根据应用于与源关联的符号的专业化函数，暂时向 shape 环境添加约束。

参数

• source (Source) – 要专业化的符号的源。

• check_fn (Callable[[Symbol], Expr]) – 一个函数，它接受一个 sympy Symbol 并返回一个 sympy 表达式，表示要应用的约束/专业化。

返回类型

Iterator[None]

produce_guards(*args, **kwargs)

类似于 produce_guards_verbose，但仅返回非 verbose 的 Python guard 表达式（不生成 verbose guards）。

返回类型

list[str]

produce_guards_expression(placeholders, *, guards=None, ignore_static=True)

预期与 evaluate_guards_expression() 一起使用。为给定的占位符生成 guards，并返回一个字符串表达式，该表达式将由 evaluate_guards_expression() 使用占位符的具体值进行求值。

返回类型

Optional[str]

produce_guards_verbose(placeholders, sources, source_ref=<function ShapeEnv.<lambda>>, *, guards=None, input_contexts=None, equalities_inputs=None, _simplified=False, ignore_static=True, langs=('python', 'verbose_python'))

生成 guards 字符串列表，当在定义了所有源的张量的上下文中进行求值时，根据 guards 列表是否为 True 来返回 True 或 False。主要由 Dynamo 使用，但这也对 guards 的手动测试很有帮助（请参阅 evaluate_guards_for_args）。

为了方便测试，源可以是一个 str，在这种情况下，我们假设它是一个 LocalSource。

simplified 允许您省略 duck sizing、equality 和 0/1 guards。这在您不关心样板 guards 的测试中很有用，并且可能也有助于用户输出（但要小心；一些 equality guards 是 nontrivial 的！如果 simplified 输出也能打印它们会很好）。它是私有的，因为它不用于正常使用。

默认返回 python guards 和带 verbose 注释的 python guards（verbose）。

返回类型

list[_ShapeGuardsHelper]

replace(expr)

将符号替换应用于给定表达式中的任何符号。

返回类型

_SympyT

set_unbacked_var_to_val(k, v)

仅在 propagate_real_tensors 时使用；注册一个未备份符号的值，该值可用作解决提示的最后手段。

simplify(expr, size_oblivious=False)

使用已知的约束和替换来简化给定的表达式。

返回类型

_SympyT

size_hint(expr, *, allow_none=False)

从底层形状获取给定表达式的大小提示。它不会引入 guard，因此仅在您可以保证代码对于任意大小仍然有效时使用（例如，用于优化决策）。

返回类型

Optional[Basic]

suppress_guards()

上下文管理器，用于忽略内部生成的所有 guards。

返回类型

_GeneratorContextManager[None]