Fix bytes_limit key of mem_stats (#243)

(cherry picked from commit c5ee54f)

Author: Ruturaj4

tests/multi_device_test.py

@@ -304,7 +304,7 @@ def test_lax_full_like_efficient(self):
       self.skipTest('Only can run test on device with mem_stats')
     mesh = Mesh(devices, axis_names=("i"))
     sharding = NamedSharding(mesh, P('i'))
-    available_memory = mem_stats['bytes_reservable_limit']
+    available_memory = mem_stats['bytes_limit']
     array_size = available_memory // (6 * len(devices)) * len(devices)
     # Set up tracemalloc to track memory usage.
     tm.start()