docstr: Transform MNIST dataset by (i) applying rotation to images, then (ii) introducing correlation between attribute (color) and label.
  Attribute (rotation angle) is independent of label Y; there is a direct-causal relationship between label Y and color.
  
  :param images: original MNIST images
  :param labels: original MNIST labels
  :param environment: Value of correlation between color and label
  :param angle: Value of rotation angle used for transforming the image
  :returns: TensorDataset containing transformed images, labels, and attributes (color, angle)
docstr_tok: ["Transform","MNIST","dataset","by","(","i",")","applying","rotation","to","images",",","then","(","ii",")","introducing","correlation","between","attribute","(","color",")","and","label",".","Attribute","(","rotation","angle",")","is","independent","of","label","Y",";","there","is","a","direct-causal","relationship","between","label","Y","and","color",".",":","param","images",":","original","MNIST","images",":","param","labels",":","original","MNIST","labels",":","param","environment",":","Value","of","correlation","between","color","and","label",":","param","angle",":","Value","of","rotation","angle","used","for","transforming","the","image",":","returns",":","TensorDataset","containing","transformed","images",",","labels",",","and","attributes","(","color",",","angle",")"]
code: def color_rot_dataset(
      self, images, labels, environment, env_id, angle
  ):
      """
      Transform MNIST dataset by (i) applying rotation to images, then (ii) introducing correlation between attribute (color) and label.
      Attribute (rotation angle) is independent of label Y; there is a direct-causal relationship between label Y and color.
  
      :param images: original MNIST images
      :param labels: original MNIST labels
      :param environment: Value of correlation between color and label
      :param angle: Value of rotation angle used for transforming the image
      :returns: TensorDataset containing transformed images, labels, and attributes (color, angle)
      """
      # Subsample 2x for computational convenience
      images = images.reshape((-1, 28, 28))[:, ::2, ::2]
      # rotate the image by angle in parameter
      images = self.rotate_dataset(images, angle)
  
      images, labels, colors = self.color_dataset(
          images, labels, environment
      )
  
      x = images  # .float().div_(255.0)
      y = labels.view(-1).long()
      angles = torch.full((y.shape[0],), env_id, dtype=torch.float32)
      a = torch.stack((colors, angles), 1)
  
      return TensorDataset(x, y, a)
  
imports: import torch
  import torchvision
  from PIL import Image
  from torch.utils.data import TensorDataset
  from torchvision import transforms
  from torchvision.datasets import MNIST
  from torchvision.transforms.functional import rotate
  from dowhy.causal_prediction.datasets.base_dataset import MultipleDomainDataset
label_desc: Use image node_id 2 for calling the MNISTCausalIndAttribute obj's underlying member method code with example usage: obj.color_rot_dataset(images, labels, environment, env_id, angle) and returns: TensorDataset