class_docstr: Tests jit targeting to different backends.
class_docstr_tok: ["Tests","jit","targeting","to","different","backends","."]
code: def testJitCpu(self):
      @partial(jax.jit, backend="cpu")
      def get_arr(scale):
          return scale + jnp.ones((2, 2))
  
      x = get_arr(0.1)
  
      a = x / x.shape[0]
      b = x + jnp.ones_like(x)
      c = x + jnp.eye(2)
  
      self.assertEqual(a.devices(), {jax.devices("cpu")[0]})
      self.assertEqual(b.devices(), {jax.devices("cpu")[0]})
      self.assertEqual(c.devices(), {jax.devices("cpu")[0]})
  
file_path: jax/tests/multibackend_test.py
imports: from functools import partial
  from absl.testing import absltest
  import numpy
  import numpy.random
  from unittest import SkipTest
  import jax
  from jax._src import test_util
  from jax import numpy
  from jax import config
label_desc: Use image node_id 5 for calling the MultiBackendTest obj's underlying member method code with example usage: obj.testJitCpu() and returns: unknown