code: def test_device_put_on_python_scalars(self):
      device = jax.devices()[0]
      int_type = dtypes.canonicalize_dtype(np.int64)
      float_type = dtypes.canonicalize_dtype(np.float64)
      complex_type = dtypes.canonicalize_dtype(np.complex128)
  
      # int
      res = np.asarray(_cpp_device_put(1, device))
      self.assertEqual(res, 1)
      self.assertEqual(res.dtype, int_type)
      # We also compare to the Python Jax API, to make sure we have the exact
      # same behavior. When Jax removes the flag and removes this feature, this
      # test will fail.
      self.assertEqual(jnp.asarray(1).dtype, res.dtype)
  
      # float
      res = np.asarray(_cpp_device_put(1.0, device))
      self.assertEqual(res, 1.0)
      self.assertEqual(res.dtype, float_type)
      self.assertEqual(jnp.asarray(1.0).dtype, res.dtype)
  
      # bool
      for bool_value in [True, False]:
          res = np.asarray(_cpp_device_put(bool_value, device))
          self.assertEqual(res, np.asarray(bool_value))
          self.assertEqual(res.dtype, np.bool_)
          self.assertEqual(jnp.asarray(bool_value).dtype, res.dtype)
  
      # Complex
      if not (
          config.enable_x64.value and jtu.test_device_matches(["tpu"])
      ):
          # No TPU support for complex128.
          res = np.asarray(_cpp_device_put(1 + 1j, device))
          self.assertEqual(res, 1 + 1j)
          self.assertEqual(res.dtype, complex_type)
          self.assertEqual(jnp.asarray(1 + 1j).dtype, res.dtype)
  
imports: import inspect
  from absl.testing import absltest
  from absl.testing import parameterized
  import jax
  from jax import dtypes
  from jax import numpy
  from jax._src import api_util
  from jax._src import config
  from jax._src import core
  from jax._src import lib
  from jax._src import test_util
  from jax._src.interpreters import pxla
  import numpy
label_desc: Use image node_id 7 for calling the JaxJitTest obj's underlying member method code with example usage: obj.test_device_put_on_python_scalars() without return types