Stores information for a maybe sliced jax.Shard. Attributes: data: The actual data of the shard. index: The index of the shard. slice_arg: Arguments for `lax.slice_in_dim` in the form of (start_idx, limit_idx, axis). If `None`, it indicates that this shard do
| 135 | |
| 136 | @dataclass |
| 137 | class _ShardInfo: |
| 138 | """Stores information for a maybe sliced jax.Shard. |
| 139 | |
| 140 | Attributes: |
| 141 | data: The actual data of the shard. |
| 142 | index: The index of the shard. |
| 143 | slice_arg: Arguments for `lax.slice_in_dim` in the form of (start_idx, limit_idx, axis). |
| 144 | If `None`, it indicates that this shard doesn't need to be sliced. |
| 145 | replica_count: The replication count for this shard. |
| 146 | """ |
| 147 | |
| 148 | data: Tensor |
| 149 | index: tuple[slice, ...] |
| 150 | slice_arg: Optional[tuple[int, int, int]] |
| 151 | replica_count: int |
| 152 | |
| 153 | def shard_coordinate(self): |
| 154 | """Gets the shard coordinate according to the zarr format used by tensorstore.""" |
| 155 | coords = [] |
| 156 | for s in self.index: |
| 157 | if s.start is None: |
| 158 | coords.append(0) |
| 159 | continue |
| 160 | size = s.stop - s.start |
| 161 | assert s.start % size == 0 |
| 162 | coords.append(s.start // size) |
| 163 | # Special case for scalar. |
| 164 | if len(coords) == 0: |
| 165 | return "0" |
| 166 | return ".".join(str(x) for x in coords) |
| 167 | |
| 168 | |
| 169 | # Tuple (and thus hashable) representation of a slice object (start, end, step). |
no outgoing calls