• Docs >
  • torch.fx.experimental
Shortcuts

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.

Docs

Access comprehensive developer documentation for PyTorch

View Docs

Tutorials

Get in-depth tutorials for beginners and advanced developers

View Tutorials

Resources

Find development resources and get your questions answered

View Resources