Project: 064_GridMix_Exploring_Spatial_Modulation_for_Neural_Fi | Function: outer_step_metagrid_same_coords (iclr)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** Spatial domain augmentation implementation that randomly subsamples training coordinates to simulate domain variations and enhance generalization across varying spatial domains

**KEY TERMS:** {
  "perm": "random permutation for sampling",
  "grid_ratio": "sampling ratio for Xsub",
  "coords_train": "Xsub coordinates",
  "feat_train": "sampled features"
}

**CONTEXT CODE:**         inner_lr,
        inner_lr_grid,
        gradient_mask,
        code_dim,
        is_train,
        gradient_checkpointing,
        loss_type,
    )
    gradient_mask = 1 - gradient_mask
    modulations = inner_loop_grad_mask_sep_lr(
        func_rep,
        modulations,
        coords_train,
        feat_train,
        inner_steps,
        inner_lr,
        inner_lr_grid,
        gradient_mask,
        code_dim,
        is_train,
        gradient_checkpointing,
        loss_type,
    )

    with torch.set_grad_enabled(is_train):
        features_recon = func_rep(coords_val, modulations)
        per_example_loss = loss_fn(features_recon, feat_val)  # features
        loss = per_example_loss.mean()

    outputs = {
        "loss": loss,
        "psnr": losses.mse2psnr(per_example_loss).mean().item(),
        "modulations": modulations,
    }

    if return_reconstructions:
        outputs["reconstructions"] = (
            features_recon[-1] if "multiscale" in loss_type else features_recon
        )

    if use_rel_loss:
        rel_loss = (
            losses.batch_mse_rel_fn(features_recon[-1], feat_val).mean()
            if "multiscale" in loss_type
            else losses.batch_mse_rel_fn(features_recon, feat_val).mean()
        )
        outputs["rel_loss"] = rel_loss

    return outputs

def outer_step_metagrid_same_coords(
    func_rep,
    coordinates,
    features,
    inner_steps,
    inner_lr,
    grid_ratio,
    is_train=False,
    return_reconstructions=False,
    gradient_checkpointing=False,
    loss_type="mse",
    modulations=0,
    use_rel_loss=False,
):
    """

    Args:
        coordinates (torch.Tensor): Shape (batch_size, *, coordinate_dim). Note this
            _must_ have a batch dimension.
        features (torch.Tensor): Shape (batch_size, *, feature_dim). Note this _must_
            have a batch dimension.
    """
    raise NotImplementedError("Function implementation has been masked for testing")


def outer_step_metagrid_part_coords(
    func_rep,
    coordinates,
    features,
    inner_steps,
    inner_lr,
    grid_ratio,
    is_train=False,
    return_reconstructions=False,
    gradient_checkpointing=False,
    loss_type="mse",
    modulations=0,
    use_rel_loss=False,
):
    # if isinstance(func_rep, DDP):
    #     func_rep = func_rep.module

    modulations = modulations.requires_grad_()
    num_points = features.shape[1]
    perm = torch.randperm(num_points)
    feat_train = features[:, perm[:int(grid_ratio*num_points)]].clone()
    coords_train = coordinates[:, perm[:int(grid_ratio*num_points)]].clone()
    
    feat_val = feat_train.clone()
    coords_val = coords_train.clone()
    # feat_val = features[:, perm[int(grid_ratio*num_points):]].clone()
    # coords_val = coordinates[:, perm[int(grid_ratio*num_points):]].clone()

    # Run inner loop
    modulations = inner_loop(
        func_rep,
        modulations,
        coords_train,
        feat_train,
        inner_steps,
        inner_lr,
        is_train,
        gradient_checkpointing,
        loss_type,
    )

    with torch.set_grad_enabled(is_train):
        features_recon = func_rep(coords_val, modulations)
        per_example_loss = loss_fn(features_recon, feat_val)  # features
        loss = per_example_loss.mean()

    outputs = {
        "loss": loss,
        "psnr": losses.mse2psnr(per_example_loss).mean().item(),
        "modulations": modulations,
    }

    if return_reconstructions:
        outputs["reconstructions"] = (
            features_recon[-1] if "multiscale" in loss_type else features_recon
        )

    if use_rel_loss:
        rel_loss = (
            losses.batch_mse_rel_fn(features_recon[-1], feat_val).mean()
            if "multiscale" in loss_type
            else losses.batch_mse_rel_fn(features_recon, feat_val).mean()
        )
        outputs["rel_loss"] = rel_loss

    return outputs


def outer_step_metagrid_part_coords(
    func_rep,
    coordinates,
    features,
    inner_steps,
    inner_lr,
    grid_ratio,


Complete the function implementation replacing "raise NotImplementedError".