code: def test_dali_sequential_tensors_to_jax_array():
      batch_size = 4
      shape = (1, 5)
  
      pipe = sequential_pipeline(batch_size, shape)
      pipe.build()
  
      for batch_id in range(100):
          # given
          dali_tensor_gpu = pipe.run()[0].as_tensor()
  
          # when
          jax_array = dax.integration._to_jax_array(dali_tensor_gpu)
  
          # then
          assert jax_array.device() == jax.devices()[0]
  
          for i in range(batch_size):
              assert jax.numpy.array_equal(
                  jax_array[i],
                  jax.numpy.full(
                      shape[
                          1:
                      ],  # TODO(awolant): Explain/fix shape consistency
                      batch_id * batch_size + i,
                      np.int32,
                  ),
              )
  
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 2 for calling a global function with example usage: test_dali_sequential_tensors_to_jax_array() without return types