More thorough propagation of host linear layout. Currently linear layout on host

can only originate from entry computation. Propagation only goes strickly down/up.
More needs to be done later if such layout can original from host compute itself.
Removed the temporary pattern match solution.

PiperOrigin-RevId: 702966364

Author: Google-ML-Automation

tests/memories_test.py

@@ -1578,7 +1578,7 @@ def test_fn(x_in, y_in):
         test_fn,
         out_shardings=(
             Layout(custom_dll, sharding),
-            Layout(custom_dll, p_sharding),
+            Layout(custom_dll_linear, p_sharding),
         ),
     )
     x_out, y_out = jit_fn(x, y)
@@ -1621,7 +1621,7 @@ def test_fn(x_in, y_in):
         test_fn,
         out_shardings=(
             Layout(custom_dll, sharding),
-            Layout(custom_dll, p_sharding),
+            Layout(custom_dll_linear, p_sharding),
         ),
     )
     x_out, y_out = jit_fn(x, y)