Optimize dynamo dynamic shape caching

[JackCaoG]

When we do

compiled_fn = torch.compile(fn, backend="openxla", dynamic=True)
# input1 and input2 are of different types
compiled_fn(input1)
compiled_fn(input2)

Dynamo will only trace the python bytecode once, but will pass input with different shapes to the same optimized_mod . The way we support the dynamic shape is to maintain a mapping between different input shapes to different xla graphs.

Currently there is an ineffiency that when we do

# when `torch.compile`d function being called
extract_internal -> extract_graph_helper -> optimized_mod

we return the optimized_mod to pytorch to replace the existing function. However in this process we did not populate the map that maps input shapes to the graph hash. We do that when we

optimized_mod (found that map is empty) -> extract_graph_helper

This is unnecessary, we should populate the mapping when we first time calling extract_graph_helper

[alanwaketan]

LGTM.