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 create_symbolic_sizes_strides_storage_offset; e.g., should they be static or dynamic.
StatelessSymbolicContext	Create symbols in create_symbolic_sizes_strides_storage_offset via a symbolic_context determination as given by DimDynamic and DimConstraint.
StatefulSymbolicContext	Create symbols in create_symbolic_sizes_strides_storage_offset via a symbolic_context determination as given by a cache of Source:Symbol.
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
is_accessor_node

torch.fx.experimental.proxy_tensor

make_fx	Given a function f, return a new function which when executed with valid arguments to f, returns an FX GraphModule representing the set of operations that were executed during the course of execution.
handle_sym_dispatch	Call into the currently active proxy tracing mode to do a SymInt/SymFloat/SymBool dispatch trace on a function that operates on these arguments.
get_proxy_mode	Current the currently active proxy tracing mode, or None if we are not currently tracing.
maybe_enable_thunkify	Within this context manager, if you are doing make_fx tracing, we will thunkify all SymNode compute and avoid tracing it into the graph unless it is actually needed.
maybe_disable_thunkify	Within a context, disable thunkification.