code: def test_jax_device_scope(self):
      import jax
      from jax.lib import xla_bridge
  
      platform = xla_bridge.get_backend().platform
  
      if platform != "gpu":
          self.skipTest("Need at least one GPU for testing")
  
      with backend.device_scope("cpu:0"):
          t = backend.numpy.ones((2, 1))
          self.assertEqual(t.device(), jax.devices("cpu")[0])
      with backend.device_scope("CPU:0"):
          t = backend.numpy.ones((2, 1))
          self.assertEqual(t.device(), jax.devices("cpu")[0])
  
      # When leaving the scope, the device should be back with gpu:0
      t = backend.numpy.ones((2, 1))
      self.assertEqual(t.device(), jax.devices("gpu")[0])
  
      # Also verify the explicit gpu device
      with backend.device_scope("gpu:0"):
          t = backend.numpy.ones((2, 1))
          self.assertEqual(t.device(), jax.devices("gpu")[0])
  
file_path: keras/keras/backend/tests/device_scope_test.py
imports: import pytest
  from keras import backend
  from keras import testing
label_desc: Use image node_id 2 for calling the DeviceTest obj's underlying member method code with example usage: obj.test_jax_device_scope() without return types