def generate_transformation_matrix(self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any]) -> Tensor:
        """Generate transformation matrices with the given input and param settings."""
        batch_prob = params["batch_prob"]
        to_apply = batch_prob > 0.5  # NOTE: in case of Relaxed Distributions.

        in_tensor = self.transform_tensor(input)
        if not to_apply.any():
            trans_matrix = self.identity_matrix(in_tensor)
        elif to_apply.all():
            trans_matrix = self.compute_transformation(in_tensor, params=params, flags=flags)
        else:
            trans_matrix_A = self.identity_matrix(in_tensor)
            trans_matrix_B = self.compute_transformation(in_tensor[to_apply], params=params, flags=flags)

            if is_autocast_enabled():
                trans_matrix_A = trans_matrix_A.type(input.dtype)
                trans_matrix_B = trans_matrix_B.type(input.dtype)

            trans_matrix = trans_matrix_A.index_put((to_apply,), trans_matrix_B)

        return trans_matrix
Generate transformation matrices with the given input and param settings.
kornia/augmentation/_2d/base.py