def _affine(self, x: torch.Tensor, a: torch.Tensor, rev: bool = False) -> tuple[Any, torch.Tensor]:
        """Perform affine coupling operation.

        Args:
            x (torch.Tensor): Input tensor (passive half)
            a (torch.Tensor): Coupling network outputs
            rev (bool, optional): Reverse the operation. Defaults to False.

        Returns:
            tuple[Any, torch.Tensor]: Transformed tensor and LogJacDet
        """
        # the entire coupling coefficient tensor is scaled down by a
        # factor of ten for stability and easier initialization.
        a *= 0.1
        ch = x.shape[1]

        sub_jac = self.clamp * torch.tanh(a[:, :ch])
        if self.GIN:
            sub_jac -= torch.mean(sub_jac, dim=self.sum_dims, keepdim=True)

        if not rev:
            return (x * torch.exp(sub_jac) + a[:, ch:], torch.sum(sub_jac, dim=self.sum_dims))

        return ((x - a[:, ch:]) * torch.exp(-sub_jac), -torch.sum(sub_jac, dim=self.sum_dims))
Perform affine coupling operation.

        Args:
            x (torch.Tensor): Input tensor (passive half)
            a (torch.Tensor): Coupling network outputs
            rev (bool, optional): Reverse the operation. Defaults to False.

        Returns:
            tuple[Any, torch.Tensor]: Transformed tensor and LogJacDet
        
src/anomalib/models/components/flow/all_in_one_block.py