Add a no input memories_test and enable memories test on vf 2x2

PiperOrigin-RevId: 641361865

Author: yashk2810

tests/memories_test.py

@@ -616,6 +616,22 @@ def _check_mem_kind(self, executable_kind, out_sharding, expected_kind):
     self.assertEqual(out_kind, expected_kind)
     self.assertEqual(executable_kind, expected_kind)
 
+  def test_compute_no_inputs(self):
+    mesh = jtu.create_global_mesh((4,), ('data'))
+
+    tpu_sharding = NamedSharding(mesh, P('data'))
+    cpu_sharding = NamedSharding(mesh, P('data'), memory_kind='pinned_host')
+
+    @functools.partial(jax.jit, out_shardings=(tpu_sharding, cpu_sharding))
+    def init():
+      tpu_array = jax.random.normal(jax.random.key(42), (16,16))
+      cpu_array = jax.random.normal(jax.random.key(42), (16,16))
+      return tpu_array, cpu_array
+
+    tpu_array, cpu_array = init()
+    self.assertEqual(tpu_array.sharding, tpu_sharding)
+    self.assertEqual(cpu_array.sharding, cpu_sharding)
+
   def test_compute_on_basic(self):
     out_s = SingleDeviceSharding(jax.devices()[0], memory_kind='pinned_host')