| 63 | |
| 64 | |
| 65 | def run_distributed_sharing_test(sharding, process_id): |
| 66 | dali_local_shard = dax.integration._to_jax_array( |
| 67 | get_dali_tensor_gpu(process_id, (1), np.int32, 0), False |
| 68 | ) |
| 69 | |
| 70 | # Note: we pass only one local shard but the array virtually |
| 71 | # combines all shards together |
| 72 | dali_sharded_array = jax.make_array_from_single_device_arrays( |
| 73 | shape=(2,), sharding=sharding, arrays=[dali_local_shard] |
| 74 | ) |
| 75 | |
| 76 | # device_buffers has been removed |
| 77 | if hasattr(dali_sharded_array, "device_buffers"): |
| 78 | # This array should be backed only by one device buffer that holds |
| 79 | # local part of the data. This buffer should be on the local device. |
| 80 | assert len(dali_sharded_array.device_buffers) == 1 |
| 81 | assert dali_sharded_array.addressable_data(0) == jnp.array([process_id]) |
| 82 | assert ( |
| 83 | dax.integration._jax_device(dali_sharded_array.addressable_data(0)) |
| 84 | == jax.local_devices()[0] |
| 85 | ) |
| 86 | assert ( |
| 87 | dax.integration._jax_device(dali_sharded_array.addressable_data(0)) |
| 88 | == jax.devices()[process_id] |
| 89 | ) |
| 90 | |
| 91 | |
| 92 | def test_positional_sharding_workflow(process_id): |