import torch


def _apply_similarity_transform_hom(
    X: torch.Tensor, R: torch.Tensor, T: torch.Tensor, s: torch.Tensor
) -> torch.Tensor:
    """
    Applies a similarity transformation parametrized with a batch of orthonormal
    matrices `R` of shape `(minibatch, d, d)`, a batch of translations `T`
    of shape `(minibatch, d)` and a batch of scaling factors `s`
    of shape `(minibatch,)` to a given `d`-dimensional poses `X`
    of shape `(minibatch, num_points, d, d)`
    Code is based on PyTorch3d.
    """
    X_tformed = X.clone()

    tform = torch.eye(4, dtype=X.dtype, device=X.device)
    tform = tform[None, :, :]
    tform = tform.repeat(X_tformed.shape[0], 1, 1)
    tform[:, :3, :3] = s[:, None, None] * R
    tform[:, :3, 3] = T
    for i in range(X_tformed.shape[0]):
        # [N,4,4] @ [4,4] = [N,4,4]
        X_tformed[i] = torch.matmul(X_tformed[i], tform[i])
    return X_tformed