(sharding)
| 225 | |
| 226 | |
| 227 | def run_sharding_test(sharding): |
| 228 | # given |
| 229 | dali_shard_0 = get_dali_tensor_gpu(0, (1), np.int32, 0) |
| 230 | dali_shard_1 = get_dali_tensor_gpu(1, (1), np.int32, 1) |
| 231 | |
| 232 | shards = [ |
| 233 | dax.integration._to_jax_array(dali_shard_0, False), |
| 234 | dax.integration._to_jax_array(dali_shard_1, False), |
| 235 | ] |
| 236 | |
| 237 | assert _jax_device(shards[0]) == jax.devices()[0] |
| 238 | assert _jax_device(shards[1]) == jax.devices()[1] |
| 239 | |
| 240 | # when |
| 241 | dali_sharded_array = jax.make_array_from_single_device_arrays( |
| 242 | shape=(2,), sharding=sharding, arrays=shards |
| 243 | ) |
| 244 | |
| 245 | # then |
| 246 | jax_sharded_array = jax.device_put(jnp.arange(2), sharding) |
| 247 | |
| 248 | assert (dali_sharded_array == jax_sharded_array).all() |
| 249 | assert len(dali_sharded_array.addressable_shards) == jax.device_count() |
| 250 | |
| 251 | assert _jax_device(dali_sharded_array.addressable_shards[0].data) == jax.devices()[0] |
| 252 | assert _jax_device(dali_sharded_array.addressable_shards[1].data) == jax.devices()[1] |
| 253 | |
| 254 | |
| 255 | def run_sharding_iterator_test(sharding): |
no test coverage detected