def METHOD_NAME():
    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._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))