🐛 Bug

When scanning over modules that doesn't contain model weights, it triggers assertion error:

 File "/workspaces/pytorch/xla/torch_xla/experimental/scan_layers.py", line 93, in scan_layers final_carry, _ = scan( File "/workspaces/pytorch/xla/torch_xla/experimental/scan.py", line 156, in scan raise ValueError(f"`xs` {xs} is an empty PyTree.") ValueError: `xs` ({}, {}) is an empty PyTree. 

To Reproduce

Run the test in 1. This test will fail and the output of fake_fa_wrapper triggers raise ValueError(f"xs {xs} is an empty PyTree.").

Expected behavior

The output of fake_fa_wrapper should not trigger assertion error.