code: def transform_pt(
      x_e: Tensor,
      translation: Tuple[float, float] = (0.0, 0.0),
      rotation: float = 0.0,
  ) -> Any:
      import torch
  
      # x_e shape: (bs, nch, x, y)
      # rotation in rad, translation in pixel
      # angles: scalar or Tensor with (bs,)
      bs = x_e.shape[0]
      theta = np.zeros((2, 3)).astype(np.float32)
      theta[0, :] = [
          np.cos(rotation),
          -np.sin(rotation),
          translation[0],
      ]
      theta[1, :] = [np.sin(rotation), np.cos(rotation), translation[1]]
      theta = np.tile(theta[None], (bs, 1, 1)).reshape(bs, 2, 3)
      # convert from pixels to relative translation, (bs, n_ch, x, y)
      theta[:, 0, 2] /= x_e.shape[2] / 2.0
      theta[:, 1, 2] /= x_e.shape[3] / 2.0
  
      # to pt
      x = x_e.raw
      theta_t = torch.tensor(theta, device=x.device)
  
      assert len(x.shape) == 4
      assert theta_t.shape[1:] == (2, 3)
  
      (
          bs,
          _,
          n_x,
          n_y,
      ) = x.shape
  
      def create_meshgrid(x: torch.Tensor) -> torch.Tensor:
          space_x = torch.linspace(-1, 1, n_x, device=x.device)
          space_y = torch.linspace(-1, 1, n_y, device=x.device)
          meshgrid = torch.meshgrid([space_x, space_y])
          ones = torch.ones(meshgrid[0].shape, device=x.device)
          gridder = torch.stack([meshgrid[1], meshgrid[0], ones], dim=2)
          grid = gridder[None, ...].repeat(bs, 1, 1, 1)[..., None]
          return grid
  
      meshgrid = create_meshgrid(x)
      theta_t = theta_t[:, None, None, :, :].repeat(1, n_x, n_y, 1, 1)
      new_coords = torch.matmul(theta_t, meshgrid)
      new_coords = new_coords.squeeze_(-1)
  
      # align_corners=True to match tf implementation
      transformed_images = torch.nn.functional.grid_sample(
          x,
          new_coords,
          mode="bilinear",
          padding_mode="zeros",
          align_corners=True,
      )
      return astensor(transformed_images)
  
imports: from typing import Tuple, Any
  import numpy
  import math
  from eagerpy import astensor, Tensor
  from eagerpy.tensor import TensorFlowTensor, PyTorchTensor
label_desc: Use image node_id 2 for calling a global function with example usage: transform_pt(x_e, translation, rotation) and returns: astensor, grid