class_docstr: Tests jit targeting to different backends.
class_docstr_tok: ["Tests","jit","targeting","to","different","backends","."]
code: def testMultiBackend(self, backend):
      if backend not in ("cpu", jtu.device_under_test(), None):
          raise SkipTest("Backend is not CPU or the device under test")
  
      @partial(jax.jit, backend=backend)
      def fun(x, y):
          return jnp.matmul(x, y)
  
      x = npr.uniform(size=(10, 10))
      y = npr.uniform(size=(10, 10))
      z_host = np.matmul(x, y)
      z = fun(x, y)
      self.assertAllClose(z, z_host, rtol=1e-2)
      correct_platform = backend if backend else jtu.device_under_test()
      self.assertEqual(list(z.devices())[0].platform, correct_platform)
  
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 1 for calling the MultiBackendTest obj's underlying member method code with example usage: obj.testMultiBackend(backend) and returns: jnp