trunk/b4eade219e7b9e1cb09e7c10ef14c25103b9ddae: [Dynamic Spec]Add ObjectSpec (#182764)

Summary

Adds ObjectSpec and extends lookup_spec_from_dynamo_source so the spec system can address tensors reached via Python attribute access (obj.weight, nn.Module params/buffers, self.x inside an instance method).

import torch

from torch.fx.experimental.dynamic_spec import (

    ObjectSpec, ParamsSpec, ShapesSpec, ShapeVar, TensorSpec,

)



class Model(torch.nn.Module):

    def __init__(self):

        super().__init__()

        self.weight = torch.nn.Parameter(torch.randn(4, 3))

    def forward(self):

        return self.weight + 1



torch.compile(

    Model(),

    shapes_spec=ShapesSpec(

        params=ParamsSpec(

            {"self": ObjectSpec({"weight": TensorSpec([ShapeVar("h"), None])})}

        )

    ),

)

What changed

torch/fx/experimental/dynamic_spec.py

• ObjectSpec — attribute-keyed container.
  
  
  
  • Constructor: ObjectSpec({name: IntermediateSpec, ...}). Values may be leaves (TensorSpec / IntVar / int / None) or another ObjectSpec for recursion.
  
  
  • Implements __getitem__ / __contains__ / __iter__ / __len__ / items() for dict-like access.
  
  
  • __repr__ matches the module's existing style (object_spec:\n .name: ...); nested children indent recursively.
  
  
  • to_jsonable() mirrors the pattern used by other spec types — recurses into spec children, passes raw leaves through.
  
  
  
  


• IntermediateSpec type alias updated to LeafSpec | ObjectSpec so containers can hold mixed leaf and nested-spec values.

torch/_dynamo/variables/builder.py

• 
  

  lookup_spec_from_dynamo_source — extended to walk Source chains:

  
  
  
  
  • LocalSource(name, is_input=True) — the root of any walk.
  
  
  • AttrSource(base, member) — descends ObjectSpec._fields[member].
  
  
  • NNModuleSource (and subclasses) — transparently unwrapped (guard-semantics marker, not an access step).
  
  
  • DictGetItemSource(UnspecializedParamBufferSource(_, '_parameters' | '_buffers'), key) — dynamo internally rewrites self.weight as self._parameters["weight"]; the walk collapses that pair into a single ("attr", key) step so the user-facing attribute name matches the spec.
  
  
  • Other source kinds return None — later container PRs (DictSpec / ListSpec) extend this dispatch.
  
  
  
  


• 
  

  wrap_module — the pre-marking loop that calls mark_static_input on each named_parameter / named_buffer now skips entries with a spec. Without this, the static stamp would be applied before wrap_tensor's spec-aware bypasses get a chance to fire, and the Parameter would still be lifted as a graph attribute.

  
  

Combined effect: nn.Parameter attributes participate in spec-driven dynamism end-to-end.

Test plan — python test/dynamo/test_dynamic_spec.py

• TestObjectSpec (data class) — empty / dict-construction / iter+items / recursive nesting / repr (none-leaf / nested-tensor / nested-objectspec) / to_jsonable.


• TestObjectSpecLookup (lookup walk, ordered simplest → most complex):
  
  
  
  • test_local_source_root_returns_top_level_spec — bare LocalSource returns the ObjectSpec itself.
  
  
  • test_attr_descends_into_objectspec — AttrSource(LocalSource(...), "weight").
  
  
  • test_nested_objectspec_walk — three-level model.inner.weight.
  
  
  • test_missing_attr_returns_none — attr not present in ObjectSpec.
  
  
  • test_attr_against_non_objectspec_returns_none — type mismatch (top-level is TensorSpec, source asks for an attr).
  
  
  • test_nn_module_source_is_unwrapped — AttrSource(NNModuleSource(LocalSource(...)), "weight").
  
  
  
  


• TestObjectSpecCompile (e2e):
  
  
  
  • test_attr_tensor_dim_dynamic — plain Python container, obj.w + 1, asserts SymInt at dim 0 and no recompile across dim-0 changes.
  
  
  • test_nn_module_parameter_dim_dynamic — realistic case: nn.Parameter reached via self.weight. Exercises NNModuleSource unwrap, UnspecializedParamBufferSource collapse, module-wrapping pre-mark skip, and wrap_tensor spec-aware bypasses.
  
  
  
  

Pull Request resolved: #182764

Approved by: https://github.com/laithsakka