def METHOD_NAME(self, n):
    try:
        if self.fake_module is not None:
            # Hacky swap. Alternatively, we could do this with overriding
            # call_module and get_attr.
            self.module = self.fake_module
        try:
            if self.fake_mode is not None:
                with self.fake_mode:
                    result = super().METHOD_NAME(n)
            else:
                result = super().METHOD_NAME(n)
        finally:
            self.module = self.real_module
        args, kwargs = self.fetch_args_kwargs_from_env(n)
        assert isinstance(args, tuple)
        assert isinstance(kwargs, dict)
        inputs = list(args) + list(kwargs.values())
        output_sharding_spec = map_aggregate(result, _extract_sharding_spec)
        n.meta["output_sharding_spec"] = output_sharding_spec
        input_sharding_spec = map_aggregate(inputs, _extract_sharding_spec)
        keys = list(n.args) + list(n.kwargs.values())
        # example:
        # code: (self.layer(inp), ) + (1, )
        # graph: add, args: ((layer, ), (1, ))
        # We need to unwrap the sharding spec here
        input_specs = dict()
        def _unwrap_sharding_spec(node_arg, sharding_spec):
            if isinstance(node_arg, Node):
                input_specs[node_arg.name] = sharding_spec
            elif isinstance(node_arg, (list, tuple)):
                for n_arg, s_spec in zip(node_arg, sharding_spec):
                    _unwrap_sharding_spec(n_arg, s_spec)
            elif isinstance(node_arg, dict):
                for n_arg, s_spec in zip(node_arg.values(), sharding_spec.values()):
                    _unwrap_sharding_spec(n_arg, s_spec)
        for i, key in enumerate(keys):
            _unwrap_sharding_spec(key, input_sharding_spec[i])
        n.meta["input_sharding_spec"] = input_specs
        return result
    except Exception as e:
        traceback.print_exc()
        raise RuntimeError(f"SpecProp error for: node={n.format_node()} with " f"meta={n.meta}") from e