code: def test_dali_tensor_gpu_to_jax_array(dtype, shape, value):
      # given
      dali_tensor_gpu = get_dali_tensor_gpu(
          value=value, shape=shape, dtype=dtype
      )
  
      # when
      jax_array = dax.integration._to_jax_array(dali_tensor_gpu)
  
      # then
      assert jax.numpy.array_equal(
          jax_array, jax.numpy.full(shape, value, dtype)
      )
  
      # Make sure JAX array is backed by the GPU
      assert jax_array.device() == jax.devices()[0]
  
file_path: DALI/dali/test/python/jax_plugin/test_integration.py
imports: import numpy
  import jax
  import jax.numpy
  import jax.dlpack
  import nvidia.dali.plugin.jax
  from nose2.tools import cartesian_params
  from utils import get_dali_tensor_gpu, sequential_pipeline
label_desc: Use image node_id 1 for calling a global function with example usage: test_dali_tensor_gpu_to_jax_array(dtype, shape, value) without return types