_produce_dynamic_shapes_for_export() — pytorch Function Reference

Source Code

benchmarks/dynamo/common.py lines 1335–1342

def _produce_dynamic_shapes_for_export(path, x):
    # mark_dynamic() is ignored for export.
    # use this to produce dynamic_shapes spec instead.
    from torch.export.dynamic_shapes import Dim

    if not isinstance(x, torch.Tensor):
        return None
    return dict.fromkeys(getattr(x, "_dynamo_dynamic_indices", {}), Dim.AUTO)