rslearn 0.0.3__py3-none-any.whl → 0.0.4__py3-none-any.whl

{rslearn-0.0.3.dist-info → rslearn-0.0.4.dist-info}/RECORD

@@ -1,7 +1,8 @@
 rslearn/__init__.py,sha256=fFmAen3vxZyosEfPbG0W46IttujYGVxzrGkJ0YutmmY,73
+rslearn/arg_parser.py,sha256=mkiZCiomUI5GNjG1jfPuTJebGHFzXbyUqe0pPwS4lTA,2055
 rslearn/const.py,sha256=FUCfsvFAs-QarEDJ0grdy0C1HjUjLpNFYGo5I2Vpc5Y,449
 rslearn/log_utils.py,sha256=unD9gShiuO7cx5Nnq8qqVQ4qrbOOwFVgcHxN5bXuiAo,941
-rslearn/main.py,sha256=vnWGvfNgj0mlUFwmOo3_OWMs2-FG8Q41LgivBLmLitA,28829
+rslearn/main.py,sha256=lAMcE4e3wCO2tVUq3bJl2oOHyztsTagtSNc0kJU7OZk,29266
 rslearn/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rslearn/config/__init__.py,sha256=Bhf2VVncdMYRC8Wfb4GsJJ13OAJYNCO_ODLSNTmBOHM,638
 rslearn/config/dataset.py,sha256=cR6Jd9ppzHgKHCteUsNapCcsJk4k5X90E71EHfbW7m0,21046
@@ -10,11 +11,11 @@ rslearn/data_sources/aws_landsat.py,sha256=GA9H04KagBDm-N37jFdh_aHCX2ZneVdnqT1SN
 rslearn/data_sources/aws_open_data.py,sha256=nU_D5cqc-wibxq4uyUNb0z-XD0Puf1gZ8v5FMiMAN5w,30258
 rslearn/data_sources/aws_sentinel1.py,sha256=cmf_ZcB7GCyFAdbwExeAwJIHqLL0JVoXtq5WcQ8UuiU,5197
 rslearn/data_sources/climate_data_store.py,sha256=Hct-0Ui-_CCQISOlzsqkK1dKz8684HRqfVUI-zXW2wA,11571
-rslearn/data_sources/copernicus.py,sha256=fIQXVDDfJwklgna6PPeCDOE6X6KSqw_XwaKcIbxJilI,36660
+rslearn/data_sources/copernicus.py,sha256=BHPyeDLeCy1-Sjyhv84snW-TnZyNVnLyn4pjjZLfTzE,36652
 rslearn/data_sources/data_source.py,sha256=69ptYhqa6pnKM04ux9hWvTPExN_lFNuU_0t_seYFnHE,3916
-rslearn/data_sources/earthdaily.py,sha256=PDsbE47mhbzzeEGiXNjZlNt-qKHYGzAjJbZTIs-Halk,18110
+rslearn/data_sources/earthdaily.py,sha256=dxOWm7Yiuh4fWVptRws_Ljh-HuNs1frf86ao91yS_80,19059
 rslearn/data_sources/earthdata_srtm.py,sha256=ysyVbVDLjhhLKdh7WKhQcwZezqvmTYaiPetTborW6zQ,11166
-rslearn/data_sources/gcp_public_data.py,sha256=O8v2iN4ym9Kwl4Zlw1FURPWnSoMD1drje8XH4xNQggE,36134
+rslearn/data_sources/gcp_public_data.py,sha256=kr9stYo7ZCvz8s4E3wmoY-yAGZoLa_9RCwjS-Q5k9dM,36128
 rslearn/data_sources/geotiff.py,sha256=sFUp919chaX4j6lQytNp__xnMLlDI3Ac3rfB6F8sgZ0,45
 rslearn/data_sources/google_earth_engine.py,sha256=hpkt74ly2lEwjRrDp8FBmGvB3MEw_mQ38Av4rQOR3_w,24246
 rslearn/data_sources/local_files.py,sha256=-XyydSPtui1m49YuP7YrNKjM5DBWMf7YgpWE9uRcvrM,18365
@@ -25,7 +26,7 @@ rslearn/data_sources/planetary_computer.py,sha256=Wchr-OmAffuVteUW6VRofIqFpE-cJq
 rslearn/data_sources/raster_source.py,sha256=b8wo55GhVLxXwx1WYLzeRAlzD_ZkE_P9tnvUOdnsfQE,689
 rslearn/data_sources/usda_cdl.py,sha256=2_V11AhPRgLEGd4U5Pmx3UvE2HWBPbsFXhUIQVRVFeE,7138
 rslearn/data_sources/usgs_landsat.py,sha256=31GmOUfmxwTE6MTiVI4psb-ciVmunuA8cfvqDuvTHPE,19312
-rslearn/data_sources/utils.py,sha256=xeLQeUh--fjnfJuyC8nZPdRxMIQdnn6VLoFlMQ32hPE,12114
+rslearn/data_sources/utils.py,sha256=oi2ybE423TLgpXlNjZ5qDQxDiwbSs7b-qD71UueQZHE,11327
 rslearn/data_sources/vector_source.py,sha256=NCa7CxIrGKe9yRT0NyyFKFQboDGDZ1h7663PV9OfMOM,44
 rslearn/data_sources/worldcereal.py,sha256=Psdf3EF3REj1WDltHWyMaICY3--KAJO_nEqpF0Gl_G8,21808
 rslearn/data_sources/worldcover.py,sha256=rimHJpQN9a56GaYxyHTOGXKzE3bkWKgd1UbH5A4aaGs,6097
@@ -59,7 +60,7 @@ rslearn/models/ssl4eo_s12.py,sha256=sOGEHcDo-rNdmEyoLu2AVEqfxRM_cv6zpfAmyn5c6tw,
 rslearn/models/swin.py,sha256=bMlGePXMFou4A_YSUZzjHgN9NniGXaCWdGQ31xHDKis,5511
 rslearn/models/task_embedding.py,sha256=Z6sf61BLCtvdrdnvjh8500b-KiFp3GeWbT4mOqpaCKk,9100
 rslearn/models/terramind.py,sha256=kipar8sMaHJJ3b8vIgL0-s4qhHcA0Vb854vmlZ9cWh4,7524
-rslearn/models/trunk.py,sha256=jm5kXlHdFkUOuIlOSwfko93Luzefkug19hKYx5l617Y,10706
+rslearn/models/trunk.py,sha256=H1QPQGAKsmocq3OiF66GW8MQI4LffupTDrgZR4Ta7QM,4708
 rslearn/models/unet.py,sha256=0xoKSsfG7y71lOqlx1F2G1H-4qq_ChjAuaAhNlTWIeo,5793
 rslearn/models/upsample.py,sha256=A0ppAFvoqSMMvESE5vxvA8giY8cToD8QoeMMPGk2tUg,965
 rslearn/models/use_croma.py,sha256=OSBqMuLp-pDtqPNWAVBfmX4wckmyYCKtUDdGCjJk_K8,17966
@@ -70,33 +71,34 @@ rslearn/models/detr/matcher.py,sha256=4h_xFlgTMEJvJ6aLZUamrKZ72L5hDk9wPglNZ81JBg
 rslearn/models/detr/position_encoding.py,sha256=8FFoBT-Jtgqk7D4qDBTbVLOeAdmjdjtJTC608TaX6yY,3869
 rslearn/models/detr/transformer.py,sha256=aK4HO7AkCZn7xGHP3Iq91w2iFPVshugOILYAjVjroCw,13971
 rslearn/models/detr/util.py,sha256=NMHhHbkIo7PoBUVbDqa2ZknJBTswmaxFCGHrPtFKnGg,676
-rslearn/models/moe/distributed.py,sha256=CFlL8eC6I4LZynz5ydcs1Xy7BuLVuEfcqHupoiKLRQ0,7948
-rslearn/models/moe/soft.py,sha256=PJgifOWBvb-ltA0NJghsOtl3fDoixOz08ZQLqtEdndU,21411
 rslearn/tile_stores/__init__.py,sha256=o_tWVKu6UwFzZbO9jn_3cmIDqc_Q3qDd6tA9If0T_Qk,2050
 rslearn/tile_stores/default.py,sha256=DEZmji2iLEVgI3abXwne6tb4C1qVtx_CaLxTpswfUV4,13852
 rslearn/tile_stores/tile_store.py,sha256=9AeYduDYPp_Ia2NMlq6osptpz_AFGIOQcLJrqZ_m-z0,10469
 rslearn/train/__init__.py,sha256=fnJyY4aHs5zQqbDKSfXsJZXY_M9fbTsf7dRYaPwZr2M,30
-rslearn/train/data_module.py,sha256=bmwMNmpNCD_6mUHoS8TGJ_6ogPD79YSaN6CWYG1Cu90,22028
-rslearn/train/dataset.py,sha256=K9ncH980rvKl08PLzzXNCd4PxZRvhSu_r7SfI0R5kAI,45217
-rslearn/train/lightning_module.py,sha256=5wzXsXf_N3r8s_qbgxZwiDj-UWkRgviuVrmbdzPZDvg,14397
+rslearn/train/data_module.py,sha256=K-nQgnOZn-KGq_G2pVOQFtWRrlWih0212i_bkXZ2bEE,23515
+rslearn/train/dataset.py,sha256=YiskNlYYcKqZxyw0Xzop1RGLbjMc-oK_rmhrSMVbTQg,51857
+rslearn/train/lightning_module.py,sha256=ge2z8trU7cMvxBeqUXC1tB44pftzitw7DRsIa6asBS4,14623
 rslearn/train/optimizer.py,sha256=EKSqkmERalDA0bF32Gey7n6z69KLyaUWKlRsGJfKBmE,927
-rslearn/train/prediction_writer.py,sha256=jiJChKmP6ZWylS3ElRVlfYBNANwYkKFp_wmTHXf5OTA,13012
-rslearn/train/scheduler.py,sha256=MBMv3TEtjEJuHZGNtP_qXKk9UenQg-JKdmbkHEV3jsc,1850
+rslearn/train/prediction_writer.py,sha256=YNs92QqPrqbREZXoE-aPa_oKQW0C9LvZAY129vyvI08,13288
+rslearn/train/scheduler.py,sha256=wFbmycMHgL6nRYeYalDjb0G8YVo8VD3T3sABS61jJ7c,2318
 rslearn/train/callbacks/__init__.py,sha256=VNV0ArZyYMvl3dGK2wl6F046khYJ1dEBlJS6G_SYNm0,47
-rslearn/train/callbacks/freeze_unfreeze.py,sha256=vuNLGyTLxyhR2Ih-GEoH-CDwFOBwBgm3yWEpn9rusBU,3766
-rslearn/train/callbacks/gradients.py,sha256=nhbA5f4QUsVfa-oV9Zj2xKk7GwWu5ejplhtIZEkRDng,4390
+rslearn/train/callbacks/adapters.py,sha256=yfv8nyCj3jmo2_dNkFrjukKxh0MHsf2xKqWwMF0QUtY,1869
+rslearn/train/callbacks/freeze_unfreeze.py,sha256=8fIzBMhCKKjpTffIeAdhdSjsBd8NjTZZEPBQaSul6Zc,17418
+rslearn/train/callbacks/gradients.py,sha256=4YqCf0tBb6E5FnyFYbveXfQFlgNPyxIXb2FCWX4-6qs,5075
 rslearn/train/callbacks/peft.py,sha256=wEOKsS3RhsRaZTXn_Kz2wdsZdIiIaZPdCJWtdJBurT8,4156
 rslearn/train/tasks/__init__.py,sha256=dag1u72x1-me6y0YcOubUo5MYZ0Tjf6-dOir9UeFNMs,75
 rslearn/train/tasks/classification.py,sha256=DI0_Wzs-9rNPWokvfxi1BIA6QyqNee42SpptQx82WHM,13182
 rslearn/train/tasks/detection.py,sha256=OoZzC8ZbmhyZ30tD-4cB-3Jj0AN6Y7hg0wk27rDguCE,22297
 rslearn/train/tasks/multi_task.py,sha256=dBWsnbvQ0CReNsbDHmZ_-sXjUE0H4S2OPcbJwMquG9g,6016
-rslearn/train/tasks/regression.py,sha256=mnlKT6zdf1fFfuOXDIAzGFZs_uU1szIoYsGuZcJJflQ,11444
-rslearn/train/tasks/segmentation.py,sha256=-CIMLYseHmuNkzrVtahILD4iMtdtN_fvW2-NcfiGq0U,20381
+rslearn/train/tasks/per_pixel_regression.py,sha256=tkVntKFzPlWFxdupPlMfhIRWlJ0UCgxg_FGhcA2-wjE,8649
+rslearn/train/tasks/regression.py,sha256=_PoxOfWNseujD4IWsuTL82fAAXgtco4WdfkNXQ68Nbg,11497
+rslearn/train/tasks/segmentation.py,sha256=xEni3CLDyetviv84XrpJg5xeJU87WHGFKTVfIeemGIY,21868
 rslearn/train/tasks/task.py,sha256=4w2xKL_U5JAtdj2dYoVv82h6xTtgUsA3IvIOcXyZecs,3887
 rslearn/train/transforms/__init__.py,sha256=BkCAzm4f-8TEhPIuyvCj7eJGh36aMkZFYlq-H_jkSvY,778
 rslearn/train/transforms/concatenate.py,sha256=sdVLJIyr9Nj2tzXEzvWFQnjJjyRSuhR_Faf6UlMIvbg,1568
 rslearn/train/transforms/crop.py,sha256=4jA3JJsC0ghicPHbfsNJ0d3WpChyvftY73ONiwQaif0,4214
 rslearn/train/transforms/flip.py,sha256=lkTeje3T8gNn2gt6957morXq1fGNho-apSpCvNp0_9o,3480
+rslearn/train/transforms/mask.py,sha256=pwt33XXWLwldLiar-PgVgBQzQd1qfL18SPz3LYQMoYM,2111
 rslearn/train/transforms/normalize.py,sha256=zYLqcDQcrjukjf5XrbFmS990PK1WQMSmHqQZKa_T040,3514
 rslearn/train/transforms/pad.py,sha256=EDswS9KYRSloM3DQlbCz6S0WYqFQJvI433qMqTtqrZw,4686
 rslearn/train/transforms/transform.py,sha256=8Q-dPrmDr0tJ9ZOwjWqWK8kbnKi4uLxEnS9Nwf6BVJk,3594
@@ -104,7 +106,7 @@ rslearn/utils/__init__.py,sha256=GNvdTUmXakiEMnLdje7k1fe5aC7SFVqP757kbpN6Fzw,558
 rslearn/utils/array.py,sha256=JwZi7o0uj-dftREzJmqrRVR2joIwBikm3Er9KeHVIZU,2402
 rslearn/utils/feature.py,sha256=lsg0WThZDJzo1mrbaL04dXYI5G3x-n5FG9aEjj7uUaI,1649
 rslearn/utils/fsspec.py,sha256=9QwN46heBhjUnth3qFeRNE3W6Wlr6dM3twYVswPnS9o,5300
-rslearn/utils/geometry.py,sha256=DTmDXNKsAaOOtm8laWELcYL6ckxBbgFJddeyGE0OMiw,15971
+rslearn/utils/geometry.py,sha256=kE4RP1g2QLcbUVF329CUSBUYHdPbpnnMuzSqLMOXvLQ,15955
 rslearn/utils/get_utm_ups_crs.py,sha256=kUrcyjCK7KWvuP1XR-nURPeRqYeRO-3L8QUJ1QTF9Ps,3599
 rslearn/utils/grid_index.py,sha256=hRmrtgpqN1pLa-djnZtgSXqKJlbgGyttGnCEmPLD0zo,2347
 rslearn/utils/jsonargparse.py,sha256=JcTKQoZ6jgwag-kSeTIEVBO9AsRj0X1oEJBsoaCazH4,658
@@ -115,9 +117,9 @@ rslearn/utils/spatial_index.py,sha256=eomJAUgzmjir8j9HZnSgQoJHwN9H0wGTjmJkMkLLfs
 rslearn/utils/sqlite_index.py,sha256=YGOJi66544e6JNtfSft6YIlHklFdSJO2duxQ4TJ2iu4,2920
 rslearn/utils/time.py,sha256=2ilSLG94_sxLP3y5RSV5L5CG8CoND_dbdzYEHVtN-I8,387
 rslearn/utils/vector_format.py,sha256=XggLCIUQBZWhOXWjvhrxBOHULpsXCbktm804DSAAink,15167
-rslearn-0.0.3.dist-info/licenses/LICENSE,sha256=_99ZWPoLdlUbqZoSC5DF4ihiNwl5rTEmBaq2fACecdg,11352
-rslearn-0.0.3.dist-info/METADATA,sha256=5aejNuiFYOCf7CNX-8TDnSSCd6f_Mddw1JO2SKN1cVE,31707
-rslearn-0.0.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-rslearn-0.0.3.dist-info/entry_points.txt,sha256=doTBQ57NT7nq-dgYGgTTw6mafcGWb_4PWYtYR4rGm50,46
-rslearn-0.0.3.dist-info/top_level.txt,sha256=XDKo90WBH8P9RQumHxo0giLJsoufT4r9odv-WE6Ahk4,8
-rslearn-0.0.3.dist-info/RECORD,,
+rslearn-0.0.4.dist-info/licenses/LICENSE,sha256=_99ZWPoLdlUbqZoSC5DF4ihiNwl5rTEmBaq2fACecdg,11352
+rslearn-0.0.4.dist-info/METADATA,sha256=_UtS9N1YTE0lKlPdei2nlGpuEHFgLiqW1h0P50EfNOU,31534
+rslearn-0.0.4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+rslearn-0.0.4.dist-info/entry_points.txt,sha256=doTBQ57NT7nq-dgYGgTTw6mafcGWb_4PWYtYR4rGm50,46
+rslearn-0.0.4.dist-info/top_level.txt,sha256=XDKo90WBH8P9RQumHxo0giLJsoufT4r9odv-WE6Ahk4,8
+rslearn-0.0.4.dist-info/RECORD,,

rslearn/models/moe/distributed.py

@@ -1,262 +0,0 @@
-"""Distributed training utilities for Soft MoE.
-
-This module provides utilities for distributed training of Soft MoE models,
-including all-gather operations and rank-based tensor splitting.
-
-Copied from
-https://raw.githubusercontent.com/lucidrains/soft-moe-pytorch/refs/heads/main/soft_moe_pytorch/distributed.py.
-"""
-
-from typing import Any
-
-import torch
-import torch.distributed as dist
-import torch.nn.functional as F
-from einops import rearrange
-from torch import Tensor, nn
-from torch.autograd import Function
-
-
-def exists(val: Any) -> bool:
-    """Check if a value exists (is not None).
-
-    Args:
-        val: The value to check.
-
-    Returns:
-        bool: True if the value is not None, False otherwise.
-    """
-    return val is not None
-
-
-def default(val: Any, d: Any) -> Any:
-    """Return the value if it exists, otherwise return the default.
-
-    Args:
-        val: The value to check.
-        d: The default value to return if val is None.
-
-    Returns:
-        Any: The value if it exists, otherwise the default.
-    """
-    return val if exists(val) else d
-
-
-def divisible_by(num: int, den: int) -> bool:
-    """Check if a number is divisible by another.
-
-    Args:
-        num: The numerator.
-        den: The denominator.
-
-    Returns:
-        bool: True if num is divisible by den, False otherwise.
-    """
-    return (num % den) == 0
-
-
-def pad_dim_to(t: Tensor, length: int, dim: int = 0) -> Tensor:
-    """Pad a tensor along a specific dimension to a target length.
-
-    Args:
-        t: The input tensor.
-        length: The target length to pad to.
-        dim: The dimension to pad along.
-
-    Returns:
-        Tensor: The padded tensor with the specified dimension padded to length.
-    """
-    pad_length = length - t.shape[dim]
-    zero_pairs = (-dim - 1) if dim < 0 else (t.ndim - dim - 1)
-    return F.pad(t, (*((0, 0) * zero_pairs), 0, pad_length))
-
-
-def all_gather_same_dim(t: Tensor) -> list[Tensor]:
-    """Gather tensors from all processes when they have the same dimension.
-
-    Args:
-        t: The tensor to gather from all processes.
-
-    Returns:
-        List[Tensor]: List of tensors gathered from all processes.
-
-    Note:
-        This function assumes all processes have tensors with the same shape.
-    """
-    world_size = dist.get_world_size()
-    t = t.contiguous()
-    gathered_tensors = [
-        torch.empty_like(t, device=t.device, dtype=t.dtype) for i in range(world_size)
-    ]
-    dist.all_gather(gathered_tensors, t)
-    return gathered_tensors
-
-
-def gather_sizes(t: Tensor, *, dim: int) -> Tensor:
-    """Gather the sizes of tensors along a specific dimension from all processes.
-
-    Args:
-        t: The input tensor.
-        dim: The dimension to gather sizes for.
-
-    Returns:
-        Tensor: Tensor containing the sizes from all processes.
-    """
-    size = torch.tensor(t.shape[dim], device=t.device, dtype=torch.long)
-    sizes = all_gather_same_dim(size)
-    return torch.stack(sizes)
-
-
-def has_only_one_value(t: Tensor) -> bool:
-    """Check if all values in a tensor are the same.
-
-    Args:
-        t: The input tensor.
-
-    Returns:
-        bool: True if all values in the tensor are identical, False otherwise.
-    """
-    return (t == t[0]).all()
-
-
-def all_gather_variable_dim(
-    t: Tensor, dim: int = 0, sizes: Tensor | None = None
-) -> tuple[Tensor, Tensor]:
-    """Gather tensors from all processes when they may have different dimensions.
-
-    Args:
-        t: The tensor to gather from all processes.
-        dim: The dimension along which tensors may vary.
-        sizes: Optional pre-computed sizes tensor. If None, will be computed.
-
-    Returns:
-        Tuple[Tensor, Tensor]: The gathered tensors and the sizes tensor.
-
-    Note:
-        This function handles the case where tensors from different processes
-        may have different sizes along the specified dimension.
-    """
-    device = t.device
-
-    if not exists(sizes):
-        sizes = gather_sizes(t, dim=dim)
-
-    if has_only_one_value(sizes):
-        gathered_tensors = all_gather_same_dim(t)
-        gathered_tensors = torch.cat(gathered_tensors, dim=dim)
-        return gathered_tensors, sizes
-
-    # Add null check for sizes
-    if sizes is None:
-        raise ValueError("sizes cannot be None")
-
-    max_size = sizes.amax().item()
-
-    padded_t = pad_dim_to(t, max_size, dim=dim)
-    gathered_tensors = all_gather_same_dim(padded_t)
-
-    gathered_tensors = torch.cat(gathered_tensors, dim=dim)
-    seq = torch.arange(max_size, device=device)
-
-    mask = rearrange(seq, "j -> 1 j") < rearrange(sizes, "i -> i 1")
-    mask = rearrange(mask, "i j -> (i j)")
-    seq = torch.arange(mask.shape[-1], device=device)
-    indices = seq[mask]
-
-    # Convert gathered_tensors to tensor before calling index_select
-    if isinstance(gathered_tensors, list):
-        gathered_tensors = torch.cat(gathered_tensors, dim=dim)
-
-    gathered_tensors = gathered_tensors.index_select(dim, indices)  # type: ignore
-
-    return gathered_tensors, sizes
-
-
-class AllGatherFunction(Function):
-    """Custom autograd function for all-gather operations.
-
-    This function provides gradient support for all-gather operations
-    by implementing custom forward and backward passes.
-    """
-
-    @staticmethod
-    def forward(
-        ctx: Any, x: Tensor, dim: int, sizes: Tensor | None
-    ) -> tuple[Tensor, Tensor]:
-        """Forward pass of the all-gather function.
-
-        Args:
-            ctx: The context object for storing information for backward pass.
-            x: The input tensor to gather.
-            dim: The dimension along which to gather.
-            sizes: Optional pre-computed sizes tensor.
-
-        Returns:
-            Tuple[Tensor, Tensor]: The gathered tensor and the sizes tensor.
-        """
-        x, batch_sizes = all_gather_variable_dim(x, dim=dim, sizes=sizes)
-        ctx.batch_sizes = batch_sizes.tolist()
-        ctx.dim = dim
-        return x, batch_sizes
-
-    @staticmethod
-    def backward(ctx: Any, grads: Tensor, _: Any) -> tuple[Tensor, None, None]:
-        """Backward pass of the all-gather function.
-
-        Args:
-            ctx: The context object containing information from forward pass.
-            grads: The gradient tensor.
-            _: Unused parameter for compatibility.
-
-        Returns:
-            Tuple[Tensor, None, None]: The gradient for the input tensor and None for other inputs.
-        """
-        batch_sizes, rank = ctx.batch_sizes, dist.get_rank()
-        grads_by_rank = grads.split(batch_sizes, dim=ctx.dim)
-        return grads_by_rank[rank], None, None
-
-
-class AllGather(nn.Module):
-    """A module that performs all-gather operations across distributed processes.
-
-    This module provides a convenient interface for gathering tensors from
-    all processes in a distributed training setup.
-    """
-
-    def __init__(self, *, dim: int = 0) -> None:
-        """Initialize the AllGather module.
-
-        Args:
-            dim: The dimension along which to gather tensors.
-        """
-        super().__init__()
-        self.dim = dim
-
-    def forward(self, x: Tensor, sizes: Tensor | None = None) -> tuple[Tensor, Tensor]:
-        """Forward pass of the all-gather operation.
-
-        Args:
-            x: The input tensor to gather from all processes.
-            sizes: Optional pre-computed sizes tensor.
-
-        Returns:
-            Tuple[Tensor, Tensor]: The gathered tensor and the sizes tensor.
-        """
-        return AllGatherFunction.apply(x, self.dim, sizes)
-
-
-def split_by_rank(x: list[Tensor]) -> Tensor:
-    """Split a list of tensors and return the tensor corresponding to the current rank.
-
-    Args:
-        x: List of tensors, one per rank.
-
-    Returns:
-        Tensor: The tensor corresponding to the current process rank.
-
-    Note:
-        This function assumes the list has one tensor per rank and returns
-        the tensor corresponding to the current process rank.
-    """
-    rank = dist.get_rank()
-    return x[rank]