jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def render_core_outside(
        self,
        rays_o,
        rays_d,
        z_vals,
        sample_dist,
        nerf,
        background_rgb=None,
        a_embedded=None,
    ):
        """
        Render background
        """
        device = rays_o.device

        batch_size, n_samples = z_vals.shape

        # Section length
        dists = z_vals[..., 1:] - z_vals[..., :-1]
        dists = torch.cat([dists, sample_dist.expand(dists[..., :1].shape)], -1)
        mid_z_vals = z_vals + dists * 0.5

        # Section midpoints
        pts = (
            rays_o[:, None, :] + rays_d[:, None, :] * mid_z_vals[..., :, None]
        )  # batch_size, n_samples, 3

        dis_to_center = torch.norm(pts, ord=2, dim=-1, keepdim=True).clip(
            1.0, 1e10
        )
        pts = torch.cat(
            [pts / dis_to_center, 1.0 / dis_to_center], dim=-1
        )  # batch_size, n_samples, 4

        dirs = rays_d[:, None, :].expand(batch_size, n_samples, 3)

        pts = pts.reshape(-1, 3 + int(self.n_outside > 0))
        dirs = dirs.reshape(-1, 3)

        a_embedded_expand = (
            a_embedded[:, None, :]
            .expand(batch_size, n_samples, -1)
            .reshape(dirs.size()[0], -1)
            if a_embedded is not None
            else None
        )

        density, sampled_color = nerf(pts, dirs, a_embedded_expand)
        alpha = 1.0 - torch.exp(
            -F.softplus(density.reshape(batch_size, n_samples)) * dists
        )
        alpha = alpha.reshape(batch_size, n_samples)
        weights = (
            alpha
            * torch.cumprod(
                torch.cat(
                    [torch.ones([batch_size, 1], device=device), 1.0 - alpha + 1e-7], -1
                ),
                -1,
            )[:, :-1]
        )
        sampled_color = sampled_color.reshape(batch_size, n_samples, 3)
        color = (weights[:, :, None] * sampled_color).sum(dim=1)
        if background_rgb is not None:
            color = color + background_rgb * (1.0 - weights.sum(dim=-1, keepdim=True))

        return {
            "color": color,
            "sampled_color": sampled_color,
            "alpha": alpha,
            "weights": weights,
        }
//  [python]
def render_core_outside(
        self,
        rays_o,
        rays_d,
        z_vals,
        sample_dist,
        nerf,
        background_rgb=None,
        a_embedded=None,
    ):
        """
        Render background
        """
        device = rays_o.device

        batch_size, n_samples = z_vals.shape

        # Section length
        dists = z_vals[..., 1:] - z_vals[..., :-1]
        dists = torch.cat([dists, sample_dist.expand(dists[..., :1].shape)], -1)
        mid_z_vals = z_vals + dists * 0.5

        # Section midpoints
        pts = (
            rays_o[:, None, :] + rays_d[:, None, :] * mid_z_vals[..., :, None]
        )  # batch_size, n_samples, 3

        dis_to_center = torch.linalg.norm(pts, ord=2, dim=-1, keepdim=True).clip(
            1.0, 1e10
        )
        pts = torch.cat(
            [pts / dis_to_center, 1.0 / dis_to_center], dim=-1
        )  # batch_size, n_samples, 4

        dirs = rays_d[:, None, :].expand(batch_size, n_samples, 3)

        pts = pts.reshape(-1, 3 + int(self.n_outside > 0))
        dirs = dirs.reshape(-1, 3)

        a_embedded_expand = (
            a_embedded[:, None, :]
            .expand(batch_size, n_samples, -1)
            .reshape(dirs.size()[0], -1)
            if a_embedded is not None
            else None
        )

        density, sampled_color = nerf(pts, dirs, a_embedded_expand)
        alpha = 1.0 - torch.exp(
            -F.softplus(density.reshape(batch_size, n_samples)) * dists
        )
        alpha = alpha.reshape(batch_size, n_samples)
        weights = (
            alpha
            * torch.cumprod(
                torch.cat(
                    [torch.ones([batch_size, 1], device=device), 1.0 - alpha + 1e-7], -1
                ),
                -1,
            )[:, :-1]
        )
        sampled_color = sampled_color.reshape(batch_size, n_samples, 3)
        color = (weights[:, :, None] * sampled_color).sum(dim=1)
        if background_rgb is not None:
            color = color + background_rgb * (1.0 - weights.sum(dim=-1, keepdim=True))

        return {
            "color": color,
            "sampled_color": sampled_color,
            "alpha": alpha,
            "weights": weights,
        }
//  [python]
def floor_loss(self, label, normals, rays_o, rays_d, rendered_depth):
        device = rays_o.device
        batch_size = label.size()[0]
        floor_mask = torch.zeros([batch_size], dtype=torch.bool)
        for f_label in self.floor_labels:
            floor_mask[get_label_id_mapping()[f_label] == label] = 1

        if torch.sum(floor_mask) == 0:
            return torch.zeros_like(normals), torch.zeros_like(normals)

        floor_normal = normals[floor_mask]

        # in gt coordinate system, z is perpendicular to road
        floor_normal_gt = torch.zeros_like(floor_normal)
        floor_normal_gt[:, 2] = 1
        floor_normal_gt = self.sfm_to_gt[:3, :3].T @ floor_normal_gt.permute(1, 0)
        floor_normal_gt = floor_normal_gt.permute(1, 0).contiguous()
        floor_normal_gt = floor_normal_gt / torch.norm(
            floor_normal_gt, dim=-1
        ).reshape(-1, 1)

        floor_error = F.l1_loss(floor_normal, floor_normal_gt, reduction="none")

        xyzs = rays_o + rays_d * rendered_depth.reshape(-1, 1)
        road_y = xyzs[floor_mask]
        y_error = torch.var(road_y)

        return floor_error, torch.ones_like(floor_error) * y_error
//  [python]
def floor_loss(self, label, normals, rays_o, rays_d, rendered_depth):
        device = rays_o.device
        batch_size = label.size()[0]
        floor_mask = torch.zeros([batch_size], dtype=torch.bool)
        for f_label in self.floor_labels:
            floor_mask[get_label_id_mapping()[f_label] == label] = 1

        if torch.sum(floor_mask) == 0:
            return torch.zeros_like(normals), torch.zeros_like(normals)

        floor_normal = normals[floor_mask]

        # in gt coordinate system, z is perpendicular to road
        floor_normal_gt = torch.zeros_like(floor_normal)
        floor_normal_gt[:, 2] = 1
        floor_normal_gt = self.sfm_to_gt[:3, :3].T @ floor_normal_gt.permute(1, 0)
        floor_normal_gt = floor_normal_gt.permute(1, 0).contiguous()
        floor_normal_gt = floor_normal_gt / torch.linalg.norm(
            floor_normal_gt, dim=-1
        ).reshape(-1, 1)

        floor_error = F.l1_loss(floor_normal, floor_normal_gt, reduction="none")

        xyzs = rays_o + rays_d * rendered_depth.reshape(-1, 1)
        road_y = xyzs[floor_mask]
        y_error = torch.var(road_y)

        return floor_error, torch.ones_like(floor_error) * y_error
//  [python]
inc= 800
n_list = np.arange(n_min,n_max,inc)

NN_loss = np.zeros([len(n_list),runs,2])
NN_weight_diff = np.zeros([len(n_list),runs,8])

for i in tqdm(range(len(n_list))):
    n = n_list[i]
    print(n)
    d = int(n*gamma1)
    h = int(n*gamma2)
    # training data with noise
    X,Y,X_test,Y_test = single_index(n,d,int(2*n/3),sigma=sigma,nonlin=sigmoid)
    for j in tqdm(range(runs)):
          W = torch.FloatTensor(d,h).normal_(mean=0,std=(1/d)**alpha1).to(device)
          b = torch.FloatTensor(h,1).normal_(mean=0,std=(1/h)**alpha2).to(device)
          s0 = torch.linalg.eigvalsh(W@torch.t(W))
          Ws = W.clone()
          bs = b.clone()
          XX = X@torch.t(X) / d
          CK0 = act(X@W)@torch.t(act(X@W)) / h
          eta_w = eta
          eta_b = eta
          XW_grad0 = grad(X@W)
          S0 = XW_grad0*((torch.ones(n, 1).to(device))@torch.t(b))
          RF0 = S0@torch.t(S0)
          NTK0 = XX * RF0
          if train_2nd:
              NTK0 = NTK0 + CK0
          s_ntk0 = torch.norm(NTK0, p=2)
          for k in range(steps):
              ys = act(X@Ws)@bs
              # update parameters
              dW = torch.t(X) @ (loss_grad(Y, ys) @ torch.t(bs) * grad(X@Ws)) / n
              Ws = Ws + eta_w * (dW - lda*Ws)

              if train_2nd:
                  db = torch.t(X@Ws)@loss_grad(Y, ys) / n
                  bs = bs + eta_b * (db - lda*bs)
          XW = act(X@Ws)
          y_NN = act(X_test@Ws)@bs
          NN_loss[i,j,0] = loss_fn(ys,Y)
          NN_loss[i,j,1] = loss_fn(y_NN,Y_test)
          s1 = torch.linalg.eigvalsh(Ws@torch.t(Ws))
          NN_weight_diff[i,j,0] = torch.abs(s1[0]-s0[0])
          NN_weight_diff[i,j,1] = torch.norm(Ws-W, p=2)
          NN_weight_diff[i,j,2] = torch.norm(Ws-W, p='fro')
          # CK model
          CK1 = act(X@Ws)@torch.t(act(X@Ws)) / h
          NN_weight_diff[i,j,3] = torch.norm(CK1-CK0, p=2)
          NN_weight_diff[i,j,4] = torch.norm(CK1-CK0, p='fro')
          # NTK model
          XW_grad = grad(X@Ws)
          S = XW_grad*((torch.ones(n, 1).to(device))@torch.t(bs))
          RF = S@torch.t(S)
          NTKs = XX * RF
          if train_2nd:
              NTKs = NTKs + CK1
          s_ntk = torch.norm(NTKs, p=2)
          NN_weight_diff[i,j,5] = torch.norm(NTKs-NTK0, p=2)
          NN_weight_diff[i,j,6] = torch.norm(NTKs-NTK0, p='fro')
          NN_weight_diff[i,j,7] = s_ntk-s_ntk0

          print("-------------------------------------------------")
          print("dimesion: %d; Training Loss: %f" % (n,loss_fn(ys,Y)))
//  [python]
inc= 800
n_list = np.arange(n_min,n_max,inc)

NN_loss = np.zeros([len(n_list),runs,2])
NN_weight_diff = np.zeros([len(n_list),runs,8])

for i in tqdm(range(len(n_list))):
    n = n_list[i]
    print(n)
    d = int(n*gamma1)
    h = int(n*gamma2)
    # training data with noise
    X,Y,X_test,Y_test = single_index(n,d,int(2*n/3),sigma=sigma,nonlin=sigmoid)
    for j in tqdm(range(runs)):
          W = torch.FloatTensor(d,h).normal_(mean=0,std=(1/d)**alpha1).to(device)
          b = torch.FloatTensor(h,1).normal_(mean=0,std=(1/h)**alpha2).to(device)
          s0 = torch.linalg.eigvalsh(W@torch.t(W))
          Ws = W.clone()
          bs = b.clone()
          XX = X@torch.t(X) / d
          CK0 = act(X@W)@torch.t(act(X@W)) / h
          eta_w = eta
          eta_b = eta
          XW_grad0 = grad(X@W)
          S0 = XW_grad0*((torch.ones(n, 1).to(device))@torch.t(b))
          RF0 = S0@torch.t(S0)
          NTK0 = XX * RF0
          if train_2nd:
              NTK0 = NTK0 + CK0
          s_ntk0 = torch.linalg.norm(NTK0, ord=2)
          for k in range(steps):
              ys = act(X@Ws)@bs
              # update parameters
              dW = torch.t(X) @ (loss_grad(Y, ys) @ torch.t(bs) * grad(X@Ws)) / n
              Ws = Ws + eta_w * (dW - lda*Ws)

              if train_2nd:
                  db = torch.t(X@Ws)@loss_grad(Y, ys) / n
                  bs = bs + eta_b * (db - lda*bs)
          XW = act(X@Ws)
          y_NN = act(X_test@Ws)@bs
          NN_loss[i,j,0] = loss_fn(ys,Y)
          NN_loss[i,j,1] = loss_fn(y_NN,Y_test)
          s1 = torch.linalg.eigvalsh(Ws@torch.t(Ws))
          NN_weight_diff[i,j,0] = torch.abs(s1[0]-s0[0])
          NN_weight_diff[i,j,1] = torch.linalg.norm(Ws-W, ord=2)
          NN_weight_diff[i,j,2] = torch.linalg.norm(Ws-W, ord='fro')
          # CK model
          CK1 = act(X@Ws)@torch.t(act(X@Ws)) / h
          NN_weight_diff[i,j,3] = torch.linalg.norm(CK1-CK0, ord=2)
          NN_weight_diff[i,j,4] = torch.linalg.norm(CK1-CK0, ord='fro')
          # NTK model
          XW_grad = grad(X@Ws)
          S = XW_grad*((torch.ones(n, 1).to(device))@torch.t(bs))
          RF = S@torch.t(S)
          NTKs = XX * RF
          if train_2nd:
              NTKs = NTKs + CK1
          s_ntk = torch.linalg.norm(NTKs, ord=2)
          NN_weight_diff[i,j,5] = torch.linalg.norm(NTKs-NTK0, ord=2)
          NN_weight_diff[i,j,6] = torch.linalg.norm(NTKs-NTK0, ord='fro')
          NN_weight_diff[i,j,7] = s_ntk-s_ntk0

          print("-------------------------------------------------")
          print("dimesion: %d; Training Loss: %f" % (n,loss_fn(ys,Y)))
//  [python]
def make_prediction(model, drs):
    norms = t.norm(drs, dim=1).reshape(-1, 1)
    energies = model(norms)

    network_derivs = vmap(jacrev(model))(norms).squeeze(-1)
    forces = -network_derivs * drs / norms
    return energies, forces
//  [python]
def scale_angle(self, scale: float, normalize_quats: bool = True) -> Rotation:
        """
        Fix the rotation axis and rescale the rotation angle.

        Args:
            scale: scale applied to the rotation angle
            normalize_quats: Whether to normalize the output quaternion
        Returns:
            A new Rotation
        """
        quats = self.get_quats()
        norm = torch.norm(quats, dim=-1)
        assert (((norm - 1).abs() < 1e-3) | (norm.abs() < 1e-3)).all(), \
            "scale_angle requires quaternions to be normalized"
        theta = torch.arccos(quats[..., 0]) # actually theta / 2
        cos_scale = torch.cos(theta * scale) / (torch.cos(theta) + 5e-8)
        sin_scale = torch.sin(theta * scale) / (torch.sin(theta) + 5e-8)
        quats_scale = torch.stack([cos_scale, sin_scale, sin_scale, sin_scale], dim=-1)
        new_quats = quats * quats_scale

        return Rotation(
            rot_mats=None,
            quats=new_quats,
            normalize_quats=normalize_quats,
        )
//  [python]
def scale_angle(self, scale: float, normalize_quats: bool = True) -> Rotation:
        """
        Fix the rotation axis and rescale the rotation angle.

        Args:
            scale: scale applied to the rotation angle
            normalize_quats: Whether to normalize the output quaternion
        Returns:
            A new Rotation
        """
        quats = self.get_quats()
        norm = torch.linalg.norm(quats, dim=-1)
        assert (((norm - 1).abs() < 1e-3) | (norm.abs() < 1e-3)).all(), \
            "scale_angle requires quaternions to be normalized"
        theta = torch.arccos(quats[..., 0]) # actually theta / 2
        cos_scale = torch.cos(theta * scale) / (torch.cos(theta) + 5e-8)
        sin_scale = torch.sin(theta * scale) / (torch.sin(theta) + 5e-8)
        quats_scale = torch.stack([cos_scale, sin_scale, sin_scale, sin_scale], dim=-1)
        new_quats = quats * quats_scale

        return Rotation(
            rot_mats=None,
            quats=new_quats,
            normalize_quats=normalize_quats,
        )
//  [python]
def knn(ref, nbr, k=1, ord=2, dim=-1, largest=False, sorted=True):
    diff = ref.unsqueeze(-2) - nbr.unsqueeze(-3)
    dist = torch.norm(diff, dim=dim, p=ord)
    return dist.topk(k, dim=dim, largest=largest, sorted=sorted)
//  [python]
def knn(ref, nbr, k=1, ord=2, dim=-1, largest=False, sorted=True):
    diff = ref.unsqueeze(-2) - nbr.unsqueeze(-3)
    dist = torch.linalg.norm(diff, dim=dim, ord=ord)
    return dist.topk(k, dim=dim, largest=largest, sorted=sorted)
//  [python]
def intersect(self, rays, distance):
        distance = distance.view(distance.shape[0], -1)

        # Min radius
        if self.min_radius > 0:
            d_from_o = torch.norm(rays[..., :3], dim=-1)
            distance_offset = distance - d_from_o.unsqueeze(1)
        else:
            distance_offset = 0

        # Add distance offset
        distance = distance + distance_offset

        return distance
//  [python]
def intersect(self, rays, distance):
        distance = distance.view(distance.shape[0], -1)

        # Min radius
        if self.min_radius > 0:
            d_from_o = torch.linalg.norm(rays[..., :3], dim=-1)
            distance_offset = distance - d_from_o.unsqueeze(1)
        else:
            distance_offset = 0

        # Add distance offset
        distance = distance + distance_offset

        return distance
//  [python]
def compute_objective(self, outputs, labels):
        """Compute the objective the FedProx client wishes to minimize."""
        current_weights = _flatten_weights_from_model(self.model, self.device)
        parameter_mu = (
            Config().clients.proximal_term_penalty_constant
            if hasattr(Config().clients, "proximal_term_penalty_constant")
            else 1
        )
        proximal_term = (
            parameter_mu
            / 2
            * torch.norm(current_weights - self.init_global_weights, p=2)
        )

        local_function = torch.nn.CrossEntropyLoss()
        function_h = local_function(outputs, labels) + proximal_term
        return function_h
//  [python]
def compute_objective(self, outputs, labels):
        """Compute the objective the FedProx client wishes to minimize."""
        current_weights = _flatten_weights_from_model(self.model, self.device)
        parameter_mu = (
            Config().clients.proximal_term_penalty_constant
            if hasattr(Config().clients, "proximal_term_penalty_constant")
            else 1
        )
        proximal_term = (
            parameter_mu
            / 2
            * torch.linalg.norm(current_weights - self.init_global_weights, ord=2)
        )

        local_function = torch.nn.CrossEntropyLoss()
        function_h = local_function(outputs, labels) + proximal_term
        return function_h
//  [python]
def __init__(self, test, device, batch_size=None, extra_args=[]):
        super().__init__(
            test=test, device=device, batch_size=batch_size, extra_args=extra_args
        )

        self.model = nn.Sequential(
            nn.Linear(1, 16),
            nn.Tanh(),
            nn.Linear(16, 16),
            nn.Tanh(),
            nn.Linear(16, 16),
            nn.Tanh(),
            nn.Linear(16, 16),
            nn.Tanh(),
            nn.Linear(16, 1),
        )
        self.model = self.model.to(device)

        r = torch.linspace(0.5, 2 * sigma, steps=self.batch_size)

        # Create a bunch of vectors that point along positive-x.
        # These are the dummy inputs to the model.
        self.drs = torch.outer(r, torch.tensor([1.0, 0, 0])).to(device=device)

        # Generate some dummy targets based off of some interpretation of the lennard_jones force.
        norms = torch.norm(self.drs, dim=1).reshape(-1, 1)
        self.norms = norms
        # Create training energies
        self.training_energies = torch.stack(list(map(lennard_jones, norms))).reshape(
            -1, 1
        )
        # Create forces with random direction vectors
        self.training_forces = torch.stack(
            [force * dr for force, dr in zip(map(lennard_jones_force, norms), self.drs)]
        )

        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=1e-3)
//  [python]
def __init__(self, test, device, batch_size=None, extra_args=[]):
        super().__init__(
            test=test, device=device, batch_size=batch_size, extra_args=extra_args
        )

        self.model = nn.Sequential(
            nn.Linear(1, 16),
            nn.Tanh(),
            nn.Linear(16, 16),
            nn.Tanh(),
            nn.Linear(16, 16),
            nn.Tanh(),
            nn.Linear(16, 16),
            nn.Tanh(),
            nn.Linear(16, 1),
        )
        self.model = self.model.to(device)

        r = torch.linspace(0.5, 2 * sigma, steps=self.batch_size)

        # Create a bunch of vectors that point along positive-x.
        # These are the dummy inputs to the model.
        self.drs = torch.outer(r, torch.tensor([1.0, 0, 0])).to(device=device)

        # Generate some dummy targets based off of some interpretation of the lennard_jones force.
        norms = torch.linalg.norm(self.drs, dim=1).reshape(-1, 1)
        self.norms = norms
        # Create training energies
        self.training_energies = torch.stack(list(map(lennard_jones, norms))).reshape(
            -1, 1
        )
        # Create forces with random direction vectors
        self.training_forces = torch.stack(
            [force * dr for force, dr in zip(map(lennard_jones_force, norms), self.drs)]
        )

        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=1e-3)
//  [python]
def boxes_dist_gpu(boxes_a, boxes_b):
    """
    :param boxes_a: (M, 7) [x, y, z, h, w, l, ry]
    :param boxes_b: (N, 7) [x, y, z, h, w, l, ry]
    :return:
        distance: (M, N)
    """
    m = len(boxes_a)
    n = len(boxes_b)
    boxes_a_corners = kitti_utils.boxes3d_to_corners3d_torch(boxes_a)
    boxes_b_corners = kitti_utils.boxes3d_to_corners3d_torch(boxes_b)
    center_distance = torch.norm(
        boxes_a[:, :3].unsqueeze(1).repeat(1, n, 1) - boxes_b[:, :3].unsqueeze(0).repeat(m, 1, 1),
        p=2, dim=-1)
    corner_distance, _ = torch.max(torch.norm(
        boxes_a_corners.view(m, 1, 8, 1, 3).repeat(1, n, 1, 8, 1)
        - boxes_b_corners.view(1, n, 1, 8, 3).repeat(m, 1, 8, 1, 1),
        p=2, dim=-1).view(m, n, 64), dim=-1)
    return 1. - center_distance / corner_distance
//  [python]
def boxes_dist_gpu(boxes_a, boxes_b):
    """
    :param boxes_a: (M, 7) [x, y, z, h, w, l, ry]
    :param boxes_b: (N, 7) [x, y, z, h, w, l, ry]
    :return:
        distance: (M, N)
    """
    m = len(boxes_a)
    n = len(boxes_b)
    boxes_a_corners = kitti_utils.boxes3d_to_corners3d_torch(boxes_a)
    boxes_b_corners = kitti_utils.boxes3d_to_corners3d_torch(boxes_b)
    center_distance = torch.linalg.norm(
        boxes_a[:, :3].unsqueeze(1).repeat(1, n, 1) - boxes_b[:, :3].unsqueeze(0).repeat(m, 1, 1),
        ord=2, dim=-1)
    corner_distance, _ = torch.max(torch.linalg.norm(
        boxes_a_corners.view(m, 1, 8, 1, 3).repeat(1, n, 1, 8, 1)
        - boxes_b_corners.view(1, n, 1, 8, 3).repeat(m, 1, 8, 1, 1),
        ord=2, dim=-1).view(m, n, 64), dim=-1)
    return 1. - center_distance / corner_distance
//  [python]
def __call__(self, inputs, targets, approximate):
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
        shape = self.feature_layer.shape

        if self.norm_function == 'neuronwise':
            H = self.hessian_function(inputs, targets, device)
            regularization = torch.zeros(1, requires_grad=True).to(device)
            for neuron_i in range(shape[0]):
                neuron_i_weights = self.feature_layer[neuron_i, :]
                for neuron_j in range(shape[0]):
                    neuron_j_weights = self.feature_layer[neuron_j, :]
                    hessian = torch.tensor(H[:, neuron_j, neuron_i, :], requires_grad=True).to(device)
                    norm_trace = torch.trace(hessian) / (1.0 * shape[1])
                    regularization = regularization + torch.dot(neuron_i_weights, neuron_j_weights) * norm_trace
        elif self.norm_function == 'layerwise_trace':
            weights_norm = torch.norm(self.feature_layer)
            regularization = weights_norm * self.hessian_function.trace(inputs, targets, approximate, device)
        elif self.norm_function == 'layerwise_eigenvalue':
            weights_norm = torch.norm(self.feature_layer)
            regularization = weights_norm * self.hessian_function.max_eigenvalue(inputs, targets, approximate, device)
        return regularization
//  [python]
def __call__(self, inputs, targets, approximate):
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
        shape = self.feature_layer.shape

        if self.norm_function == 'neuronwise':
            H = self.hessian_function(inputs, targets, device)
            regularization = torch.zeros(1, requires_grad=True).to(device)
            for neuron_i in range(shape[0]):
                neuron_i_weights = self.feature_layer[neuron_i, :]
                for neuron_j in range(shape[0]):
                    neuron_j_weights = self.feature_layer[neuron_j, :]
                    hessian = torch.tensor(H[:, neuron_j, neuron_i, :], requires_grad=True).to(device)
                    norm_trace = torch.trace(hessian) / (1.0 * shape[1])
                    regularization = regularization + torch.dot(neuron_i_weights, neuron_j_weights) * norm_trace
        elif self.norm_function == 'layerwise_trace':
            weights_norm = torch.linalg.norm(self.feature_layer)
            regularization = weights_norm * self.hessian_function.trace(inputs, targets, approximate, device)
        elif self.norm_function == 'layerwise_eigenvalue':
            weights_norm = torch.linalg.matrix_norm(self.feature_layer, dim=tuple(range(len(shape))))
            regularization = weights_norm * self.hessian_function.max_eigenvalue(inputs, targets, approximate, device)
        return regularization
//  [python]
def evaluate_R_t_tensor(R_gt, t_gt, R, t, q_gt=None):
    t = t.flatten()
    t_gt = t_gt.flatten().to(t.device)

    eps = torch.tensor(1e-8).to(t.device)
    err_q = torch.arccos(torch.max(torch.min((torch.trace(R @ R_gt.transpose(0, 1)) - 1) * 0.5,
                                 torch.tensor(1.0, device=R.device)), torch.tensor(-1.0, device=R.device)))

    t_gt_ = t_gt / (t_gt.norm() + eps)
    loss_t = torch.max(eps, 1.0 - torch.sum(t * t_gt_)**2)#torch.clamp((), min=eps)
    err_t = torch.arccos(torch.sqrt(1 - loss_t+eps))

    if torch.sum(torch.isnan(err_q)) or torch.sum(torch.isnan(err_t)):
        print("This should never happen! Debug here", R_gt, t_gt, R, t, q_gt)
        import IPython
        IPython.embed()

    return err_q, err_t
//  [python]
def evaluate_R_t_tensor(R_gt, t_gt, R, t, q_gt=None):
    t = t.flatten()
    t_gt = t_gt.flatten().to(t.device)

    eps = torch.tensor(1e-8).to(t.device)
    err_q = torch.arccos(torch.max(torch.min((torch.trace(R @ R_gt.transpose(0, 1)) - 1) * 0.5,
                                 torch.tensor(1.0, device=R.device)), torch.tensor(-1.0, device=R.device)))

    t_gt_ = t_gt / (torch.linalg.norm(t_gt) + eps)
    loss_t = torch.max(eps, 1.0 - torch.sum(t * t_gt_)**2)#torch.clamp((), min=eps)
    err_t = torch.arccos(torch.sqrt(1 - loss_t+eps))

    if torch.sum(torch.isnan(err_q)) or torch.sum(torch.isnan(err_t)):
        print("This should never happen! Debug here", R_gt, t_gt, R, t, q_gt)
        import IPython
        IPython.embed()

    return err_q, err_t
//  [python]
def residuIL_loss(self, targets, predictions, pred_residuals):

        # train policy (detach discriminator)
        target_prediction_differences = targets - predictions

        # bc_regularization is optionally used according to Appendix of the
        # paper (Tabel 4 and 5)
        policy_loss = (
            2 * (target_prediction_differences) * pred_residuals.detach()
        ).mean(-1) + self._bc_regulatization_weight * (
            torch.square(target_prediction_differences)).mean(-1)

        # train discriminator (detach policy)
        discriminator_loss = -(
            2 * (target_prediction_differences).detach() * pred_residuals -
            pred_residuals * pred_residuals).mean(-1)

        # f_norm_penalty is used according to Appendix of the paper (Tabel 4 and 5)
        discriminator_loss = (
            discriminator_loss +
            self._f_norm_penalty_weight * torch.norm(pred_residuals))

        return policy_loss, discriminator_loss
//  [python]
def residuIL_loss(self, targets, predictions, pred_residuals):

        # train policy (detach discriminator)
        target_prediction_differences = targets - predictions

        # bc_regularization is optionally used according to Appendix of the
        # paper (Tabel 4 and 5)
        policy_loss = (
            2 * (target_prediction_differences) * pred_residuals.detach()
        ).mean(-1) + self._bc_regulatization_weight * (
            torch.square(target_prediction_differences)).mean(-1)

        # train discriminator (detach policy)
        discriminator_loss = -(
            2 * (target_prediction_differences).detach() * pred_residuals -
            pred_residuals * pred_residuals).mean(-1)

        # f_norm_penalty is used according to Appendix of the paper (Tabel 4 and 5)
        discriminator_loss = (
            discriminator_loss +
            self._f_norm_penalty_weight * torch.linalg.norm(pred_residuals))

        return policy_loss, discriminator_loss
//  [python]
def evaluate_R_t_tensor_batch(R_gt, t_gt, R, t, q_gt=None):
    t = t.flatten(-2, -1)
    t_gt = t_gt.flatten().to(t.device)

    eps = torch.tensor(1e-15).to(t.device)
    # rotation error
    err_q = torch.arccos(torch.clamp((torch.diagonal(R @ R_gt.repeat(R.shape[0], 1, 1).transpose(-2, -1), dim1=-2, dim2=-1).sum(-1) - 1) * 0.5, min= -1.0, max=1.0))
    # translation error
    t = t / (torch.norm(t, dim=-1) + eps).unsqueeze(-1)
    t_gt_ = t_gt / (torch.norm(t_gt) + eps)
    loss_t = torch.clamp((1.0 - torch.sum(t * t_gt_, dim=-1)**2), min=eps)
    err_t = torch.arccos(torch.sqrt(1 - loss_t + 1e-8))

    if torch.sum(torch.isnan(err_q)) or torch.sum(torch.isnan(err_t)):
        # This should never happen! Debug here
        print(R_gt, t_gt, R, t, q_gt)
        import IPython
        IPython.embed()

    return err_q, err_t
//  [python]
def evaluate_R_t_tensor_batch(R_gt, t_gt, R, t, q_gt=None):
    t = t.flatten(-2, -1)
    t_gt = t_gt.flatten().to(t.device)

    eps = torch.tensor(1e-15).to(t.device)
    # rotation error
    err_q = torch.arccos(torch.clamp((torch.diagonal(R @ R_gt.repeat(R.shape[0], 1, 1).transpose(-2, -1), dim1=-2, dim2=-1).sum(-1) - 1) * 0.5, min= -1.0, max=1.0))
    # translation error
    t = t / (torch.norm(t, dim=-1) + eps).unsqueeze(-1)
    t_gt_ = t_gt / (torch.linalg.norm(t_gt) + eps)
    loss_t = torch.clamp((1.0 - torch.sum(t * t_gt_, dim=-1)**2), min=eps)
    err_t = torch.arccos(torch.sqrt(1 - loss_t + 1e-8))

    if torch.sum(torch.isnan(err_q)) or torch.sum(torch.isnan(err_t)):
        # This should never happen! Debug here
        print(R_gt, t_gt, R, t, q_gt)
        import IPython
        IPython.embed()

    return err_q, err_t
//  [python]
def random_get_rays_at(self, img_idx, batch_size, pose = None):
        pose_cur = self.get_pose(img_idx, pose)
        
        pixels_x = torch.randint(low=0, high=self.W, size=[batch_size]).cpu()
        pixels_y = torch.randint(low=0, high=self.H, size=[batch_size]).cpu()
            
        color = self.images[img_idx][(pixels_y, pixels_x)]    # batch_size, 3
        mask = self.masks[img_idx][(pixels_y, pixels_x)][:,None]     # batch_size, 3
        pts_target = self.pts[img_idx][(pixels_y*pixels_x)]
        p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1).float()  # batch_size, 3
        p = p.to(self.intrinsics_all_inv.device) 
        self.intrinsics_all_inv = self.intrinsics_all_inv.to(p.device)
        p = torch.matmul(self.intrinsics_all_inv[img_idx, None, :3, :3], p[:, :, None]).squeeze() # batch_size, 3
        rays_v = p / torch.norm(p, p=2, dim=-1, keepdim=True)    # batch_size, 3
        rays_v = torch.matmul(pose_cur[None, :3, :3], rays_v[:, :, None]).squeeze()  # batch_size, 3
        rays_o = pose_cur[None, :3, 3].expand(rays_v.shape) # batch_size, 3
        
        normal_sample = None
        if self.use_normal:
            normal_sample = self.normals[img_idx][(pixels_y, pixels_x)].cuda()

        planes_sample = None
        if self.use_planes:
            planes_sample = self.planes[img_idx][(pixels_y, pixels_x)].unsqueeze(-1).cuda()
        
        subplanes_sample = None
        if self.use_plane_offset_loss:
            subplanes_sample = self.subplanes[img_idx][(pixels_y, pixels_x)].unsqueeze(-1).cuda()

        return torch.cat([rays_o.cpu(), rays_v.cpu(), color, mask, pts_target], dim=-1).cuda(), pixels_x, pixels_y, normal_sample, planes_sample, subplanes_sample
//  [python]
def random_get_rays_at(self, img_idx, batch_size, pose = None):
        pose_cur = self.get_pose(img_idx, pose)
        
        pixels_x = torch.randint(low=0, high=self.W, size=[batch_size]).cpu()
        pixels_y = torch.randint(low=0, high=self.H, size=[batch_size]).cpu()
            
        color = self.images[img_idx][(pixels_y, pixels_x)]    # batch_size, 3
        mask = self.masks[img_idx][(pixels_y, pixels_x)][:,None]     # batch_size, 3
        pts_target = self.pts[img_idx][(pixels_y*pixels_x)]
        p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1).float()  # batch_size, 3
        p = p.to(self.intrinsics_all_inv.device) 
        self.intrinsics_all_inv = self.intrinsics_all_inv.to(p.device)
        p = torch.matmul(self.intrinsics_all_inv[img_idx, None, :3, :3], p[:, :, None]).squeeze() # batch_size, 3
        rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True)    # batch_size, 3
        rays_v = torch.matmul(pose_cur[None, :3, :3], rays_v[:, :, None]).squeeze()  # batch_size, 3
        rays_o = pose_cur[None, :3, 3].expand(rays_v.shape) # batch_size, 3
        
        normal_sample = None
        if self.use_normal:
            normal_sample = self.normals[img_idx][(pixels_y, pixels_x)].cuda()

        planes_sample = None
        if self.use_planes:
            planes_sample = self.planes[img_idx][(pixels_y, pixels_x)].unsqueeze(-1).cuda()
        
        subplanes_sample = None
        if self.use_plane_offset_loss:
            subplanes_sample = self.subplanes[img_idx][(pixels_y, pixels_x)].unsqueeze(-1).cuda()

        return torch.cat([rays_o.cpu(), rays_v.cpu(), color, mask, pts_target], dim=-1).cuda(), pixels_x, pixels_y, normal_sample, planes_sample, subplanes_sample
//  [python]
def frob_normalize(x: torch.Tensor):
    slice_shape = x.shape[-2:]
    # THIS VIEW FAILS WHEN  C_out/C_in dims are SIMPLY broadcast into existance from a single batch. Need own memory?
    x_slices = x.view(-1, *slice_shape).detach()
    norms = torch.norm(x_slices, ord='fro', dim=(1, 2)).view(-1, 1)

    # first put into stack of slices, then into stack of rows
    x_row_form = x.view(-1, *slice_shape).view(x_slices.shape[0], -1)
    # first transpose to get shapes correct, second to bring back to og form
    x_row_form_normed = torch.div(x_row_form, norms)

    return x_row_form_normed.view(x.shape)
//  [python]
def frob_normalize(x: torch.Tensor):
    slice_shape = x.shape[-2:]
    # THIS VIEW FAILS WHEN  C_out/C_in dims are SIMPLY broadcast into existance from a single batch. Need own memory?
    x_slices = x.view(-1, *slice_shape).detach()
    norms = torch.linalg.norm(x_slices, ord='fro', dim=(1, 2)).view(-1, 1)

    # first put into stack of slices, then into stack of rows
    x_row_form = x.view(-1, *slice_shape).view(x_slices.shape[0], -1)
    # first transpose to get shapes correct, second to bring back to og form
    x_row_form_normed = torch.div(x_row_form, norms)

    return x_row_form_normed.view(x.shape)
//  [python]
def latentReg(z, reg):
    latentreg = 0
    for latent_i in z:
        latentreg += 1 / reg * torch.norm(latent_i)
    return latentreg
//  [python]
def latentReg(z, reg):
    latentreg = 0
    for latent_i in z:
        latentreg += 1 / reg * torch.linalg.norm(latent_i)
    return latentreg
//  [python]
def l2_norm(a, b):
    return torch.norm(a-b,  p = 2, dim = -1)
//  [python]
def l2_norm(a, b):
    return torch.linalg.norm(a-b,  ord = 2, dim = -1)
//  [python]
def update_model( self ):
        dataset_train = self.build_dataset()
        self._debug_mode = False
        if self._debug_mode and os.path.exists(self._feature_cache):
            print("using feature cache")
            with open(self._feature_cache, 'rb') as f:
                train_data = torch.load(f)
        else:
            print("generating bbox features")
            train_data = self.get_features(dataset_train)
            saved = torch.save(train_data, self._feature_cache)

        # normalize the training features. This may require some experimentation
        # depending on how the model handles normalization

        train_data = torch.norm(train_data, dim=1, p=self._norm_degree)
        
        # Train ReMax model and then save model
        self.remax_model = ReMax(train_data)

        self.save_final_model()
//  [python]
def update_model( self ):
        dataset_train = self.build_dataset()
        self._debug_mode = False
        if self._debug_mode and os.path.exists(self._feature_cache):
            print("using feature cache")
            with open(self._feature_cache, 'rb') as f:
                train_data = torch.load(f)
        else:
            print("generating bbox features")
            train_data = self.get_features(dataset_train)
            saved = torch.save(train_data, self._feature_cache)

        # normalize the training features. This may require some experimentation
        # depending on how the model handles normalization

        train_data = torch.linalg.norm(train_data, dim=1, ord=self._norm_degree)
        
        # Train ReMax model and then save model
        self.remax_model = ReMax(train_data)

        self.save_final_model()
//  [python]
def measure(self, wrapped_model: Models.ForwardHookedOutput, dataset: DatasetWrapper, shared_cache=None) -> pd.DataFrame:
        if shared_cache is None:
            shared_cache = SharedMeasurementVarsCache()

        wrapped_model.base_model.eval()
        device = next(wrapped_model.parameters()).device

        class_means, class_num_samples = shared_cache.get_train_class_means_nums(wrapped_model, dataset)
        global_mean = shared_cache.calc_global_mean(class_means, class_num_samples)
        # classwise_cov_within = shared_cache.get_train_class_covariance(wrapped_model, dataset)

        out: List[Dict[str, Any]] = []
        for layer_name in tqdm.tqdm(class_means.keys(), desc='  ETFAngle, calculating', leave=False):

            layer_rel_class_means = (class_means[layer_name] - global_mean[layer_name]).flatten(start_dim=1).to('cpu')
            # layer_classwise_cov_within = classwise_cov_within[layer_name].to('cpu')
            # layer_cov_between = torch.matmul(layer_rel_class_means.T, layer_rel_class_means) / dataset.num_classes
            # layer_class_means = class_means[layer_name]

            layer_rel_class_means_norms = torch.norm(layer_rel_class_means, dim=1)  # Class means in rows
            normed_layer_rel_class_means = torch.nn.functional.normalize(layer_rel_class_means, dim=1)  # Class means in rows

            class_means_cos_angles = normed_layer_rel_class_means @ normed_layer_rel_class_means.T
            class_means_cos_angles_plus_correction = class_means_cos_angles + 1 / (dataset.num_classes - 1)

            for l_class_idx in range(dataset.num_classes):
                out.append({'value': layer_rel_class_means_norms[l_class_idx].item(), 'layer_name': layer_name,
                            'l_ord': l_class_idx, 'r_ord': l_class_idx, 'type': 'norm'})
                for r_class_idx in range(dataset.num_classes):
                    if l_class_idx == r_class_idx:
                        continue
                    out.append({'value': class_means_cos_angles_plus_correction[l_class_idx][r_class_idx].item(),
                                'layer_name': layer_name,
                                'l_ord': l_class_idx, 'r_ord': r_class_idx, 'type': 'angle'})

        return pd.DataFrame(out)
//  [python]
def measure(self, wrapped_model: Models.ForwardHookedOutput, dataset: DatasetWrapper, shared_cache=None) -> pd.DataFrame:
        if shared_cache is None:
            shared_cache = SharedMeasurementVarsCache()

        wrapped_model.base_model.eval()
        device = next(wrapped_model.parameters()).device

        class_means, class_num_samples = shared_cache.get_train_class_means_nums(wrapped_model, dataset)
        global_mean = shared_cache.calc_global_mean(class_means, class_num_samples)
        # classwise_cov_within = shared_cache.get_train_class_covariance(wrapped_model, dataset)

        out: List[Dict[str, Any]] = []
        for layer_name in tqdm.tqdm(class_means.keys(), desc='  ETFAngle, calculating', leave=False):

            layer_rel_class_means = (class_means[layer_name] - global_mean[layer_name]).flatten(start_dim=1).to('cpu')
            # layer_classwise_cov_within = classwise_cov_within[layer_name].to('cpu')
            # layer_cov_between = torch.matmul(layer_rel_class_means.T, layer_rel_class_means) / dataset.num_classes
            # layer_class_means = class_means[layer_name]

            layer_rel_class_means_norms = torch.linalg.norm(layer_rel_class_means, dim=1)  # Class means in rows
            normed_layer_rel_class_means = torch.nn.functional.normalize(layer_rel_class_means, dim=1)  # Class means in rows

            class_means_cos_angles = normed_layer_rel_class_means @ normed_layer_rel_class_means.T
            class_means_cos_angles_plus_correction = class_means_cos_angles + 1 / (dataset.num_classes - 1)

            for l_class_idx in range(dataset.num_classes):
                out.append({'value': layer_rel_class_means_norms[l_class_idx].item(), 'layer_name': layer_name,
                            'l_ord': l_class_idx, 'r_ord': l_class_idx, 'type': 'norm'})
                for r_class_idx in range(dataset.num_classes):
                    if l_class_idx == r_class_idx:
                        continue
                    out.append({'value': class_means_cos_angles_plus_correction[l_class_idx][r_class_idx].item(),
                                'layer_name': layer_name,
                                'l_ord': l_class_idx, 'r_ord': r_class_idx, 'type': 'angle'})

        return pd.DataFrame(out)
//  [python]
def normalize_vector(v, dim, eps=1e-6):
    return v / (torch.norm(v, p=2, dim=dim, keepdim=True) + eps)
//  [python]
def normalize_vector(v, dim, eps=1e-6):
    return v / (torch.linalg.norm(v, ord=2, dim=dim, keepdim=True) + eps)
//  [python]
def measure(self, wrapped_model: Models.ForwardHookedOutput, dataset: DatasetWrapper, shared_cache=None) -> pd.DataFrame:
        if shared_cache is None:
            shared_cache = SharedMeasurementVarsCache()

        wrapped_model.base_model.eval()
        device = next(wrapped_model.parameters()).device

        # class_means, class_num_samples = shared_cache.get_train_class_means_nums(wrapped_model, dataset)
        class_means, class_num_samples = shared_cache.get_test_class_means_nums(wrapped_model, dataset)
        global_mean = shared_cache.calc_global_mean(class_means, class_num_samples)
        # classwise_cov_within = shared_cache.get_train_class_covariance(wrapped_model, dataset)

        out: List[Dict[str, Any]] = []
        for layer_name in tqdm.tqdm(class_means.keys(), desc='  ETFAngle, calculating', leave=False):

            layer_rel_class_means = (class_means[layer_name] - global_mean[layer_name]).flatten(start_dim=1).to('cpu')
            # layer_classwise_cov_within = classwise_cov_within[layer_name].to('cpu')
            # layer_cov_between = torch.matmul(layer_rel_class_means.T, layer_rel_class_means) / dataset.num_classes
            # layer_class_means = class_means[layer_name]

            layer_rel_class_means_norms = torch.norm(layer_rel_class_means, dim=1)  # Class means in rows
            normed_layer_rel_class_means = torch.nn.functional.normalize(layer_rel_class_means, dim=1)  # Class means in rows

            class_means_cos_angles = normed_layer_rel_class_means @ normed_layer_rel_class_means.T
            class_means_cos_angles_plus_correction = class_means_cos_angles + 1 / (dataset.num_classes - 1)

            for l_class_idx in range(dataset.num_classes):
                out.append({'value': layer_rel_class_means_norms[l_class_idx].item(), 'layer_name': layer_name,
                            'l_ord': l_class_idx, 'r_ord': l_class_idx, 'type': 'norm'})
                for r_class_idx in range(dataset.num_classes):
                    if l_class_idx == r_class_idx:
                        continue
                    out.append({'value': class_means_cos_angles_plus_correction[l_class_idx][r_class_idx].item(),
                                'layer_name': layer_name,
                                'l_ord': l_class_idx, 'r_ord': r_class_idx, 'type': 'angle'})

        return pd.DataFrame(out)
//  [python]
def measure(self, wrapped_model: Models.ForwardHookedOutput, dataset: DatasetWrapper, shared_cache=None) -> pd.DataFrame:
        if shared_cache is None:
            shared_cache = SharedMeasurementVarsCache()

        wrapped_model.base_model.eval()
        device = next(wrapped_model.parameters()).device

        # class_means, class_num_samples = shared_cache.get_train_class_means_nums(wrapped_model, dataset)
        class_means, class_num_samples = shared_cache.get_test_class_means_nums(wrapped_model, dataset)
        global_mean = shared_cache.calc_global_mean(class_means, class_num_samples)
        # classwise_cov_within = shared_cache.get_train_class_covariance(wrapped_model, dataset)

        out: List[Dict[str, Any]] = []
        for layer_name in tqdm.tqdm(class_means.keys(), desc='  ETFAngle, calculating', leave=False):

            layer_rel_class_means = (class_means[layer_name] - global_mean[layer_name]).flatten(start_dim=1).to('cpu')
            # layer_classwise_cov_within = classwise_cov_within[layer_name].to('cpu')
            # layer_cov_between = torch.matmul(layer_rel_class_means.T, layer_rel_class_means) / dataset.num_classes
            # layer_class_means = class_means[layer_name]

            layer_rel_class_means_norms = torch.linalg.norm(layer_rel_class_means, dim=1)  # Class means in rows
            normed_layer_rel_class_means = torch.nn.functional.normalize(layer_rel_class_means, dim=1)  # Class means in rows

            class_means_cos_angles = normed_layer_rel_class_means @ normed_layer_rel_class_means.T
            class_means_cos_angles_plus_correction = class_means_cos_angles + 1 / (dataset.num_classes - 1)

            for l_class_idx in range(dataset.num_classes):
                out.append({'value': layer_rel_class_means_norms[l_class_idx].item(), 'layer_name': layer_name,
                            'l_ord': l_class_idx, 'r_ord': l_class_idx, 'type': 'norm'})
                for r_class_idx in range(dataset.num_classes):
                    if l_class_idx == r_class_idx:
                        continue
                    out.append({'value': class_means_cos_angles_plus_correction[l_class_idx][r_class_idx].item(),
                                'layer_name': layer_name,
                                'l_ord': l_class_idx, 'r_ord': r_class_idx, 'type': 'angle'})

        return pd.DataFrame(out)
//  [python]
def gen_random_rays_at(self, img_idx, batch_size):
        """
        Generate random rays at world space from one camera.
        """
        pixels_x = torch.randint(low=0, high=self.W, size=[batch_size])
        pixels_y = torch.randint(low=0, high=self.H, size=[batch_size])
        color = self.images[img_idx][(pixels_y, pixels_x)]    # batch_size, 3
        mask = self.masks[img_idx][(pixels_y, pixels_x)]      # batch_size, 3
        p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1).float()  # batch_size, 3
        p = torch.matmul(self.intrinsics_all_inv[img_idx, None, :3, :3], p[:, :, None]).squeeze()
        rays_v = p / torch.norm(p, p=2, dim=-1, keepdim=True)    # batch_size, 3
        rays_v = torch.matmul(self.pose_all[img_idx, None, :3, :3], rays_v[:, :, None]).squeeze()  # batch_size, 3
        rays_o = self.pose_all[img_idx, None, :3, 3].expand(rays_v.shape) # batch_size, 3

        return torch.cat([rays_o, rays_v, color, mask[:, :1]], dim=-1)
//  [python]
def gen_random_rays_at(self, img_idx, batch_size):
        """
        Generate random rays at world space from one camera.
        """
        pixels_x = torch.randint(low=0, high=self.W, size=[batch_size])
        pixels_y = torch.randint(low=0, high=self.H, size=[batch_size])
        color = self.images[img_idx][(pixels_y, pixels_x)]    # batch_size, 3
        mask = self.masks[img_idx][(pixels_y, pixels_x)]      # batch_size, 3
        p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1).float()  # batch_size, 3
        p = torch.matmul(self.intrinsics_all_inv[img_idx, None, :3, :3], p[:, :, None]).squeeze()
        rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True)    # batch_size, 3
        rays_v = torch.matmul(self.pose_all[img_idx, None, :3, :3], rays_v[:, :, None]).squeeze()  # batch_size, 3
        rays_o = self.pose_all[img_idx, None, :3, 3].expand(rays_v.shape) # batch_size, 3

        return torch.cat([rays_o, rays_v, color, mask[:, :1]], dim=-1)