Improve clarity for flattened graph inputs

[GregoryComer]

🚀 The feature, motivation and pitch

When passing complex model inputs, such as (nested) dictionaries, it can be a little unclear on the order of flattened inputs. This can be resolved by looking at the exported IR or program dump, but these aren't easy to parse without detailed knowledge of export/ET internals.

If such a thing doesn't already exist, it would be nice to provide a way for users to easily see what the expected input format is. Maybe printing the exported function signature? This could maybe be incorporated into the higher-level export API. Or as a verbose logging option when lowering? Implementation needs thinking through. Ideally, it would be easy for a user to understand what order to pass inputs at runtime without needing to dig into verbose, low-level info.

Alternatives

No response

Additional context

No response

RFC (Optional)

No response

cc @JacobSzwejbka @angelayi @mergennachin @byjlw

[JacobSzwejbka]

Hmm I'm not sure I agree with its important for users to understand 'how' their thing was flattened vs that it was flattened at all. Ideally we would just make it obvious when a user needs to call tree flatten before passing their input to the runtime, or ideally just detect it and flatten it ourselves in things like the module layer.

[GregoryComer]

The main problem I want to solve is to just make it easy to know what the actual runtime inputs should be. I know we have the pytree extension, but it's a bit of extra work to use, and from what I've seen, people often just want to pass the tensors in directly. But they may not know what order they should be in, especially when there are dictionaries involved. This is motivated by an internal question, where I realized it doesn't seem like there is any good way to actually figure out what the flattened inputs look like without dumping IR.