code: def test_dali_sequential_sharded_tensors_to_jax_sharded_array_iterator_multiple_outputs():
      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,
          multiple_outputs=True,
      )
  
      pipe_1 = sequential_sharded_pipeline(
          batch_size=batch_size,
          shape=shape,
          device_id=1,
          shard_id=1,
          shard_size=batch_size,
          multiple_outputs=True,
      )
  
      output_names = ["data_0", "data_1", "data_2"]
  
      # when
      dali_iterator = DALIGenericIterator(
          [pipe_0, pipe_1], output_names, size=batch_size * 10
      )
  
      for batch_id, batch in enumerate(dali_iterator):
          # then
          # check values for all outputs
          # for the data_0 values should be the same as in the single output example
          # for data_1 values are the same + 0.25, for data_2 the same + 0.5
          for output_id, output_name in enumerate(output_names):
              jax_array = batch[output_name]
  
              assert jax.numpy.array_equal(
                  jax_array.device_buffers[0],
                  jax.numpy.stack(
                      [
                          jax.numpy.full(
                              shape[1:],
                              value + output_id * 0.25,
                              np.float32,
                          )
                          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 + output_id * 0.25,
                              np.float32,
                          )
                          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.devices()[0]
              )
              assert (
                  jax_array.device_buffers[1].device()
                  == jax.devices()[1]
              )
  
      # Assert correct number of batches returned from the iterator
      assert batch_id == 4
  
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 3 for calling a global function with example usage: test_dali_sequential_sharded_tensors_to_jax_sharded_array_iterator_multiple_outputs() without return types