torch.fx.experimental¶
Warning
These APIs are experimental and subject to change without notice.
torch.fx.experimental.symbolic_shapes¶
torch.fx.experimental.symbolic_shapes
provides interfaces for interacting with
our symbolic shapes reasoning system that is used heavily in torch.compile. Although
this is not generally considered public API, when writing framework code in PyTorch
as well as extensions to PyTorch (e.g., in custom operator implementations), you may
need to make use of these APIs to setup dynamic shapes support appropriately.
ShapeEnv |
|
DimDynamic |
Controls how to perform symbol allocation for a dimension. |
StrictMinMaxConstraint |
For clients: the size at this dimension must be within 'vr' (which specifies a lower and upper bound, inclusive-inclusive) AND it must be non-negative and should not be 0 or 1 (but see NB below). |
RelaxedUnspecConstraint |
For clients: no explicit constraint; constraint is whatever is implicitly inferred by guards from tracing. |
EqualityConstraint |
Represent and decide various kinds of equality constraints between input sources. |
SymbolicContext |
Data structure specifying how we should create symbols in |
StatelessSymbolicContext |
Create symbols in |
StatefulSymbolicContext |
Create symbols in |
SubclassSymbolicContext |
The correct symbolic context for a given inner tensor of a traceable tensor subclass may differ from that of the outer symbolic context. |
DimConstraints |
Custom solver for a system of constraints on symbolic dimensions. |
ShapeEnvSettings |
Encapsulates all shape env settings that could potentially affect FakeTensor dispatch. |
ConvertIntKey |
|
CallMethodKey |
|
PropagateUnbackedSymInts |
|
DivideByKey |
|
InnerTensorKey |
|
hint_int |
Retrieve the hint for an int (based on the underlying real values as observed at runtime). |
is_concrete_int |
Utility to check if underlying object in SymInt is concrete value. |
is_concrete_bool |
Utility to check if underlying object in SymBool is concrete value. |
has_free_symbols |
Faster version of bool(free_symbols(val)) |
definitely_true |
Returns True only if we can tell that a is True, possibly introducing a guard in the process. |
definitely_false |
Returns True only if we can tell that a is False, possibly introducing a guard in the process. |
guard_size_oblivious |
Perform a guard on a symbolic boolean expression in a size oblivious way. |
parallel_or |
Evaluate the logical OR of several arguments, avoiding guarding on unbacked SymInts if another argument is definitely True. |
parallel_and |
Evaluate the logical FALSE of several arguments, avoiding guarding on unbacked SymInts if another argument is definitely False. |
sym_eq |
Like ==, but when run on list/tuple, it will recursively test equality and use sym_and to join the results together, without guarding. |
constrain_range |
Applies a constraint that the passed in SymInt must lie between min-max inclusive-inclusive, WITHOUT introducing a guard on the SymInt (meaning that it can be used on unbacked SymInts). |
constrain_unify |
Given two SymInts, constrain them so that they must be equal. |
canonicalize_bool_expr |
Canonicalize a boolean expression by transforming it into a lt / le inequality and moving all the non-constant terms to the rhs. |
statically_known_true |
Returns True if x can be simplified to a constant and is true. |
lru_cache |
|
check_consistent |
Test that two "meta" values (typically either Tensor or SymInt) have the same values, e.g., after retracing. |
compute_unbacked_bindings |
After having run fake tensor propagation and producing example_value result, traverse example_value looking for freshly bound unbacked symbols and record their paths for later. |
rebind_unbacked |
Suppose we are retracing a pre-existing FX graph that previously had fake tensor propagation (and therefore unbacked SymInts). |
resolve_unbacked_bindings |