perf(torch): inline GLOBAL_STATE_TRACKER getattr in get_device()/distribution() hot paths

keras/src/backend/torch/core.py

@@ -72,7 +72,7 @@ def device_scope(device_name):
 
 
 def get_device():
-    device = global_state.get_global_attribute("torch_device", None)
+    device = getattr(global_state.GLOBAL_STATE_TRACKER, "torch_device", None)
     if device is None:
         return DEFAULT_DEVICE
     return device

keras/src/backend/torch/core_test.py

@@ -5,7 +5,10 @@
 
 from keras.src import backend
 from keras.src import testing
+from keras.src.backend.torch.core import DEFAULT_DEVICE
 from keras.src.backend.torch.core import convert_to_tensor
+from keras.src.backend.torch.core import device_scope
+from keras.src.backend.torch.core import get_device
 from keras.src.backend.torch.core import slice as torch_slice
 
 
@@ -104,3 +107,11 @@ def test_slice_fast_path_accepts_symint(self):
         shape = [batch, 2, 2]
         result = torch_slice(x, start_indices, shape)
         self.assertEqual(tuple(result.shape), (2, 2, 2))
+
+    def test_get_device_reflects_device_scope(self):
+        """get_device() returns the default device outside any scope and the
+        scoped device inside device_scope."""
+        self.assertEqual(get_device(), DEFAULT_DEVICE)
+        with device_scope("cpu:0"):
+            self.assertEqual(get_device(), "cpu:0")
+        self.assertEqual(get_device(), DEFAULT_DEVICE)

keras/src/distribution/distribution_lib.py

@@ -815,7 +815,9 @@ def distribute_tensor(tensor, layout):
 @keras_export("keras.distribution.distribution")
 def distribution():
     """Retrieve the current distribution from global context."""
-    return global_state.get_global_attribute(GLOBAL_ATTRIBUTE_NAME)
+    return getattr(
+        global_state.GLOBAL_STATE_TRACKER, GLOBAL_ATTRIBUTE_NAME, None
+    )
 
 
 @keras_export("keras.distribution.set_distribution")