code: def test_dali_sequential_sharded_tensors_to_jax_sharded_array_manuall():
      assert (
          jax.device_count() > 1
      ), "Multigpu test requires more than one GPU"
  
      # given
      pipe_0 = sequential_sharded_pipeline(
          batch_size=batch_size,
          shape=shape,
          device_id=0,
          shard_id=0,
          shard_size=batch_size,
      )
      pipe_0.build()
  
      pipe_1 = sequential_sharded_pipeline(
          batch_size=batch_size,
          shape=shape,
          device_id=1,
          shard_id=1,
          shard_size=batch_size,
      )
      pipe_1.build()
  
      for batch_id in range(100):
          dali_tensor_gpu_0 = pipe_0.run()[0].as_tensor()
          dali_tensor_gpu_1 = pipe_1.run()[0].as_tensor()
  
          jax_shard_0 = dax.integration._to_jax_array(dali_tensor_gpu_0)
          jax_shard_1 = dax.integration._to_jax_array(dali_tensor_gpu_1)
  
          assert jax_shard_0.device() == jax.devices()[0]
          assert jax_shard_1.device() == jax.devices()[1]
  
          # when
          jax_array = jax.device_put_sharded(
              [jax_shard_0, jax_shard_1],
              [jax_shard_0.device(), jax_shard_1.device()],
          )
  
          # then
          # Assert that all values are as expected
          # In first iteration, first shard should be:
          # [[0 0 0 0 0]
          #  [1 1 1 1 1]
          #  [2 2 2 2 2]
          #  [3 3 3 3 3]]
          # And second shard should be:
          # [[4 4 4 4 4]
          #  [5 5 5 5 5]
          #  [6 6 6 6 6]
          #  [7 7 7 7 7]]
          # Then, in second iteration first shard should be:
          # [[4 4 4 4 4]
          #  [5 5 5 5 5]
          #  [6 6 6 6 6]
          #  [7 7 7 7 7]]
          # And second shard should be:
          # [[ 8  8  8  8  8]
          #  [ 9  9  9  9  9]
          #  [10 10 10 10 10]
          #  [11 11 11 11 11]]
          assert jax.numpy.array_equal(
              jax_array.device_buffers[0],
              jax.numpy.stack(
                  [
                      jax.numpy.full(shape[1:], value, np.int32)
                      for value in range(
                          batch_id * batch_size,
                          (batch_id + 1) * batch_size,
                      )
                  ]
              ),
          )
          assert jax.numpy.array_equal(
              jax_array.device_buffers[1],
              jax.numpy.stack(
                  [
                      jax.numpy.full(shape[1:], value, np.int32)
                      for value in range(
                          (batch_id + 1) * batch_size,
                          (batch_id + 2) * batch_size,
                      )
                  ]
              ),
          )
  
          # Assert correct backing devices for shards
          assert (
              jax_array.device_buffers[0].device()
              == jax_shard_0.device()
          )
          assert (
              jax_array.device_buffers[1].device()
              == jax_shard_1.device()
          )
  
imports: import jax
  import numpy
  import nvidia.dali.plugin.jax
  from nvidia.dali import pipeline_def
  import nvidia.dali.fn
  import nvidia.dali.types
  from nvidia.dali.plugin.jax import DALIGenericIterator, data_iterator
  from nvidia.dali.plugin.base_iterator import LastBatchPolicy
  from jax.sharding import PositionalSharding, NamedSharding, PartitionSpec, Mesh
  from jax.experimental import mesh_utils
  from utils import get_dali_tensor_gpu, iterator_function_def
  import jax.numpy
  import itertools
  from nose_utils import raises
label_desc: Use image node_id 2 for calling a global function with example usage: test_dali_sequential_sharded_tensors_to_jax_sharded_array_manuall() without return types