jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def forward(self, x: Tensor, v: Tensor) -> Tuple[Tensor, Tensor]:
        r"""Applies a gated equivariant operation to node features and vector
        features.

        Args:
            x (torch.Tensor): The scalar features of the nodes.
            v (torch.Tensor): The vector features of the nodes.
        """
        vec1 = torch.linalg.norm(self.vec1_proj(v), dim=-2)
        vec2 = self.vec2_proj(v)

        x = torch.cat([x, vec1], dim=-1)
        x, v = torch.split(self.update_net(x), self.out_channels, dim=-1)
        v = v.unsqueeze(1) * vec2

        if self.act is not None:
            x = self.act(x)

        return x, v
//  [python]
def max_min_norm(self, vec: Tensor) -> Tensor:
        r"""Applies max-min normalization to the input tensor.

        .. math::
            \text{dist} = ||\text{vec}||_2
            \text{direct} = \frac{\text{vec}}{\text{dist}}
            \text{max\_val} = \max(\text{dist})
            \text{min\_val} = \min(\text{dist})
            \text{delta} = \text{max\_val} - \text{min\_val}
            \text{dist} = \frac{\text{dist} - \text{min\_val}}{\text{delta}}
            \text{normed\_vec} = \max(0, \text{dist}) \cdot \text{direct}

        Args:
            vec (torch.Tensor): The input tensor.
        """
        dist = torch.norm(vec, dim=1, keepdim=True)

        if (dist == 0).all():
            return torch.zeros_like(vec)

        dist = dist.clamp(min=self.eps)
        direct = vec / dist

        max_val, _ = dist.max(dim=-1)
        min_val, _ = dist.min(dim=-1)
        delta = (max_val - min_val).view(-1)
        delta = torch.where(delta == 0, torch.ones_like(delta), delta)
        dist = (dist - min_val.view(-1, 1, 1)) / delta.view(-1, 1, 1)

        return dist.relu() * direct
//  [python]
def max_min_norm(self, vec: Tensor) -> Tensor:
        r"""Applies max-min normalization to the input tensor.

        .. math::
            \text{dist} = ||\text{vec}||_2
            \text{direct} = \frac{\text{vec}}{\text{dist}}
            \text{max\_val} = \max(\text{dist})
            \text{min\_val} = \min(\text{dist})
            \text{delta} = \text{max\_val} - \text{min\_val}
            \text{dist} = \frac{\text{dist} - \text{min\_val}}{\text{delta}}
            \text{normed\_vec} = \max(0, \text{dist}) \cdot \text{direct}

        Args:
            vec (torch.Tensor): The input tensor.
        """
        dist = torch.linalg.norm(vec, dim=1, keepdim=True)

        if (dist == 0).all():
            return torch.zeros_like(vec)

        dist = dist.clamp(min=self.eps)
        direct = vec / dist

        max_val, _ = dist.max(dim=-1)
        min_val, _ = dist.min(dim=-1)
        delta = (max_val - min_val).view(-1)
        delta = torch.where(delta == 0, torch.ones_like(delta), delta)
        dist = (dist - min_val.view(-1, 1, 1)) / delta.view(-1, 1, 1)

        return dist.relu() * direct
//  [python]
def rigid_from_3_points(N, Ca, C, non_ideal=False, eps=1e-8):
    #N, Ca, C - [B,L, 3]
    #R - [B,L, 3, 3], det(R)=1, inv(R) = R.T, R is a rotation matrix
    B,L = N.shape[:2]
    
    v1 = C-Ca
    v2 = N-Ca
    e1 = v1/(torch.norm(v1, dim=-1, keepdim=True)+eps)
    u2 = v2-(torch.einsum('bli, bli -> bl', e1, v2)[...,None]*e1)
    e2 = u2/(torch.norm(u2, dim=-1, keepdim=True)+eps)
    e3 = torch.cross(e1, e2, dim=-1)
    R = torch.cat([e1[...,None], e2[...,None], e3[...,None]], axis=-1) #[B,L,3,3] - rotation matrix
    
    if non_ideal:
        v2 = v2/(torch.norm(v2, dim=-1, keepdim=True)+eps)
        cosref = torch.sum(e1*v2, dim=-1) # cosine of current N-CA-C bond angle
        costgt = cos_ideal_NCAC.item()
        cos2del = torch.clamp( cosref*costgt + torch.sqrt((1-cosref*cosref)*(1-costgt*costgt)+eps), min=-1.0, max=1.0 )
        cosdel = torch.sqrt(0.5*(1+cos2del)+eps)
        sindel = torch.sign(costgt-cosref) * torch.sqrt(1-0.5*(1+cos2del)+eps)
        Rp = torch.eye(3, device=N.device).repeat(B,L,1,1)
        Rp[:,:,0,0] = cosdel
        Rp[:,:,0,1] = -sindel
        Rp[:,:,1,0] = sindel
        Rp[:,:,1,1] = cosdel
    
        R = torch.einsum('blij,bljk->blik', R,Rp)

    return R, Ca
//  [python]
def rigid_from_3_points(N, Ca, C, non_ideal=False, eps=1e-8):
    #N, Ca, C - [B,L, 3]
    #R - [B,L, 3, 3], det(R)=1, inv(R) = R.T, R is a rotation matrix
    B,L = N.shape[:2]
    
    v1 = C-Ca
    v2 = N-Ca
    e1 = v1/(torch.linalg.norm(v1, dim=-1, keepdim=True)+eps)
    u2 = v2-(torch.einsum('bli, bli -> bl', e1, v2)[...,None]*e1)
    e2 = u2/(torch.linalg.norm(u2, dim=-1, keepdim=True)+eps)
    e3 = torch.cross(e1, e2, dim=-1)
    R = torch.cat([e1[...,None], e2[...,None], e3[...,None]], axis=-1) #[B,L,3,3] - rotation matrix
    
    if non_ideal:
        v2 = v2/(torch.linalg.norm(v2, dim=-1, keepdim=True)+eps)
        cosref = torch.sum(e1*v2, dim=-1) # cosine of current N-CA-C bond angle
        costgt = cos_ideal_NCAC.item()
        cos2del = torch.clamp( cosref*costgt + torch.sqrt((1-cosref*cosref)*(1-costgt*costgt)+eps), min=-1.0, max=1.0 )
        cosdel = torch.sqrt(0.5*(1+cos2del)+eps)
        sindel = torch.sign(costgt-cosref) * torch.sqrt(1-0.5*(1+cos2del)+eps)
        Rp = torch.eye(3, device=N.device).repeat(B,L,1,1)
        Rp[:,:,0,0] = cosdel
        Rp[:,:,0,1] = -sindel
        Rp[:,:,1,0] = sindel
        Rp[:,:,1,1] = cosdel
    
        R = torch.einsum('blij,bljk->blik', R,Rp)

    return R, Ca
//  [python]
def test_ema_fp32(self):
        dtype = torch.float

        model = DummyModule().to(dtype)
        optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
        state = deepcopy(model.state_dict())
        config = EMAConfig(ema_fp32=True)
        ema = EMA(model, config)

        x = torch.randn(32)
        y = model(x.to(dtype))
        loss = y.sum()
        loss.backward()
        optimizer.step()

        ema.step(model)

        for key, param in model.state_dict().items():
            prev_param = state[key]
            ema_param = ema.get_model().state_dict()[key]

            if "version" in key:
                # Do not decay a model.version pytorch param
                continue
            self.assertIn(key, ema.fp32_params)

            # EMA update is done in fp32, and hence the EMA param must be
            # closer to the EMA update done in fp32 than in fp16.
            self.assertLessEqual(
                torch.norm(
                    ema_param.float()
                    - (
                        config.ema_decay * prev_param.float()
                        + (1 - config.ema_decay) * param.float()
                    )
                    .to(dtype)
                    .float()
                ),
                torch.norm(
                    ema_param.float()
                    - (
                        config.ema_decay * prev_param + (1 - config.ema_decay) * param
                    ).float()
                ),
            )
            self.assertTorchAllClose(
                ema_param,
                (
                    config.ema_decay * prev_param.float()
                    + (1 - config.ema_decay) * param.float()
                ).to(dtype),
            )
//  [python]
def test_ema_fp32(self):
        dtype = torch.float

        model = DummyModule().to(dtype)
        optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
        state = deepcopy(model.state_dict())
        config = EMAConfig(ema_fp32=True)
        ema = EMA(model, config)

        x = torch.randn(32)
        y = model(x.to(dtype))
        loss = y.sum()
        loss.backward()
        optimizer.step()

        ema.step(model)

        for key, param in model.state_dict().items():
            prev_param = state[key]
            ema_param = ema.get_model().state_dict()[key]

            if "version" in key:
                # Do not decay a model.version pytorch param
                continue
            self.assertIn(key, ema.fp32_params)

            # EMA update is done in fp32, and hence the EMA param must be
            # closer to the EMA update done in fp32 than in fp16.
            self.assertLessEqual(
                torch.linalg.norm(
                    ema_param.float()
                    - (
                        config.ema_decay * prev_param.float()
                        + (1 - config.ema_decay) * param.float()
                    )
                    .to(dtype)
                    .float()
                ),
                torch.linalg.norm(
                    ema_param.float()
                    - (
                        config.ema_decay * prev_param + (1 - config.ema_decay) * param
                    ).float()
                ),
            )
            self.assertTorchAllClose(
                ema_param,
                (
                    config.ema_decay * prev_param.float()
                    + (1 - config.ema_decay) * param.float()
                ).to(dtype),
            )
//  [python]
def forward(
        self,
        pos: Tensor,
        batch: Tensor,
    ) -> Tuple[Tensor, Tensor, Tensor]:
        r"""Computes the pairwise distances between atoms in the molecule.

        Args:
            pos (torch.Tensor): The positions of the atoms in the molecule.
            batch (torch.Tensor): A batch vector, which assigns each node to a
                specific example.

        Returns:
            edge_index (torch.Tensor): The indices of the edges in the graph.
            edge_weight (torch.Tensor): The distances between connected nodes.
            edge_vec (torch.Tensor): The vector differences between connected
                nodes.
        """
        edge_index = radius_graph(
            pos,
            r=self.cutoff,
            batch=batch,
            loop=self.add_self_loops,
            max_num_neighbors=self.max_num_neighbors,
        )
        edge_vec = pos[edge_index[0]] - pos[edge_index[1]]

        if self.add_self_loops:
            mask = edge_index[0] != edge_index[1]
            edge_weight = torch.zeros(edge_vec.size(0), device=edge_vec.device)
            edge_weight[mask] = torch.norm(edge_vec[mask], dim=-1)
        else:
            edge_weight = torch.norm(edge_vec, dim=-1)

        return edge_index, edge_weight, edge_vec
//  [python]
def forward(
        self,
        pos: Tensor,
        batch: Tensor,
    ) -> Tuple[Tensor, Tensor, Tensor]:
        r"""Computes the pairwise distances between atoms in the molecule.

        Args:
            pos (torch.Tensor): The positions of the atoms in the molecule.
            batch (torch.Tensor): A batch vector, which assigns each node to a
                specific example.

        Returns:
            edge_index (torch.Tensor): The indices of the edges in the graph.
            edge_weight (torch.Tensor): The distances between connected nodes.
            edge_vec (torch.Tensor): The vector differences between connected
                nodes.
        """
        edge_index = radius_graph(
            pos,
            r=self.cutoff,
            batch=batch,
            loop=self.add_self_loops,
            max_num_neighbors=self.max_num_neighbors,
        )
        edge_vec = pos[edge_index[0]] - pos[edge_index[1]]

        if self.add_self_loops:
            mask = edge_index[0] != edge_index[1]
            edge_weight = torch.zeros(edge_vec.size(0), device=edge_vec.device)
            edge_weight[mask] = torch.linalg.norm(edge_vec[mask], dim=-1)
        else:
            edge_weight = torch.linalg.norm(edge_vec, dim=-1)

        return edge_index, edge_weight, edge_vec
//  [python]
def forward(
        self,
        z: Tensor,
        pos: Tensor,
        batch: Tensor,
    ) -> Tuple[Tensor, Tensor]:
        r"""Computes the scalar and vector features of the nodes.

        Args:
            z (torch.Tensor): The atomic numbers.
            pos (torch.Tensor): The coordinates of the atoms.
            batch (torch.Tensor): A batch vector, which assigns each node to a
                specific example.

        Returns:
            x (torch.Tensor): The scalar features of the nodes.
            vec (torch.Tensor): The vector features of the nodes.
        """
        x = self.embedding(z)
        edge_index, edge_weight, edge_vec = self.distance(pos, batch)
        edge_attr = self.distance_expansion(edge_weight)
        mask = edge_index[0] != edge_index[1]
        edge_vec[mask] = edge_vec[mask] / torch.norm(edge_vec[mask],
                                                     dim=1).unsqueeze(1)
        edge_vec = self.sphere(edge_vec)
        x = self.neighbor_embedding(z, x, edge_index, edge_weight, edge_attr)
        vec = torch.zeros(x.size(0), ((self.lmax + 1)**2) - 1, x.size(1),
                          dtype=x.dtype, device=x.device)
        edge_attr = self.edge_embedding(edge_index, edge_attr, x)

        for attn in self.vis_mp_layers[:-1]:
            dx, dvec, dedge_attr = attn(x, vec, edge_index, edge_weight,
                                        edge_attr, edge_vec)
            x = x + dx
            vec = vec + dvec
            edge_attr = edge_attr + dedge_attr

        dx, dvec, _ = self.vis_mp_layers[-1](x, vec, edge_index, edge_weight,
                                             edge_attr, edge_vec)
        x = x + dx
        vec = vec + dvec

        x = self.out_norm(x)
        vec = self.vec_out_norm(vec)

        return x, vec
//  [python]
def forward(
        self,
        z: Tensor,
        pos: Tensor,
        batch: Tensor,
    ) -> Tuple[Tensor, Tensor]:
        r"""Computes the scalar and vector features of the nodes.

        Args:
            z (torch.Tensor): The atomic numbers.
            pos (torch.Tensor): The coordinates of the atoms.
            batch (torch.Tensor): A batch vector, which assigns each node to a
                specific example.

        Returns:
            x (torch.Tensor): The scalar features of the nodes.
            vec (torch.Tensor): The vector features of the nodes.
        """
        x = self.embedding(z)
        edge_index, edge_weight, edge_vec = self.distance(pos, batch)
        edge_attr = self.distance_expansion(edge_weight)
        mask = edge_index[0] != edge_index[1]
        edge_vec[mask] = edge_vec[mask] / torch.linalg.norm(edge_vec[mask],
                                                     dim=1).unsqueeze(1)
        edge_vec = self.sphere(edge_vec)
        x = self.neighbor_embedding(z, x, edge_index, edge_weight, edge_attr)
        vec = torch.zeros(x.size(0), ((self.lmax + 1)**2) - 1, x.size(1),
                          dtype=x.dtype, device=x.device)
        edge_attr = self.edge_embedding(edge_index, edge_attr, x)

        for attn in self.vis_mp_layers[:-1]:
            dx, dvec, dedge_attr = attn(x, vec, edge_index, edge_weight,
                                        edge_attr, edge_vec)
            x = x + dx
            vec = vec + dvec
            edge_attr = edge_attr + dedge_attr

        dx, dvec, _ = self.vis_mp_layers[-1](x, vec, edge_index, edge_weight,
                                             edge_attr, edge_vec)
        x = x + dx
        vec = vec + dvec

        x = self.out_norm(x)
        vec = self.vec_out_norm(vec)

        return x, vec
//  [python]
def find_max_distance(points):
    # 获取点的数量
    n_points = points.size(0)

    # 初始化最大距离为0
    max_distance = 0

    # 遍历每一对点
    for i in range(n_points):
        for j in range(i + 1, n_points):
            # 计算两点之间的欧几里得距离
            distance = torch.norm(points[i] - points[j])

            # 更新最大距离
            max_distance = max(max_distance, distance)

    # 返回最大距离
    return max_distance
//  [python]
def find_max_distance(points):
    # 获取点的数量
    n_points = points.size(0)

    # 初始化最大距离为0
    max_distance = 0

    # 遍历每一对点
    for i in range(n_points):
        for j in range(i + 1, n_points):
            # 计算两点之间的欧几里得距离
            distance = torch.linalg.norm(points[i] - points[j])

            # 更新最大距离
            max_distance = max(max_distance, distance)

    # 返回最大距离
    return max_distance
//  [python]
def parse_floor_plane(floor_plane):
    """
    Takes floor plane in the optimization form (Bx3 with a,b,c * d) and parses into
    (a,b,c,d) from with (a,b,c) normal facing "up in the camera frame and d the offset.
    """
    floor_offset = torch.norm(floor_plane, dim=-1, keepdim=True)
    floor_normal = floor_plane / floor_offset

    # in camera system -y is up, so floor plane normal y component should never be positive
    #       (assuming the camera is not sideways or upside down)
    neg_mask = floor_normal[..., 1:2] > 0.0
    floor_normal = torch.where(
        neg_mask.expand_as(floor_normal), -floor_normal, floor_normal
    )
    floor_offset = torch.where(neg_mask, -floor_offset, floor_offset)
    floor_plane_4d = torch.cat([floor_normal, floor_offset], dim=-1)

    return floor_plane_4d
//  [python]
def parse_floor_plane(floor_plane):
    """
    Takes floor plane in the optimization form (Bx3 with a,b,c * d) and parses into
    (a,b,c,d) from with (a,b,c) normal facing "up in the camera frame and d the offset.
    """
    floor_offset = torch.linalg.norm(floor_plane, dim=-1, keepdim=True)
    floor_normal = floor_plane / floor_offset

    # in camera system -y is up, so floor plane normal y component should never be positive
    #       (assuming the camera is not sideways or upside down)
    neg_mask = floor_normal[..., 1:2] > 0.0
    floor_normal = torch.where(
        neg_mask.expand_as(floor_normal), -floor_normal, floor_normal
    )
    floor_offset = torch.where(neg_mask, -floor_offset, floor_offset)
    floor_plane_4d = torch.cat([floor_normal, floor_offset], dim=-1)

    return floor_plane_4d
//  [python]
def gradFlow(self, path):
        '''
        A method to monitor the flow of gradient across the layers. Use it to monitor exploding and
        vanishing gradients. ``scaleRho`` must be tweaked to ensure proper gradient flow. Usually
        monitoring it for first few epochs is good enough.

        Usage:

        .. code-block:: python
            
            net = Network(netParams)
            net.gradFlow(path_to_save)
        '''
        gradNorm = lambda x: torch.norm(x).item()/torch.numel(x)
        grad = []

        for l in self.blocks:
            # print(l)
            if l.gradLog is True:
                if l.weightNorm is True:
                    grad.append(gradNorm(l.weightOp.weight_g.grad))
                else:
                    grad.append(gradNorm(l.weightOp.weight.grad))

        plt.figure()
        plt.semilogy(grad)
        plt.savefig(path + 'gradFlow.png')
        plt.close()
//  [python]
def gradFlow(self, path):
        '''
        A method to monitor the flow of gradient across the layers. Use it to monitor exploding and
        vanishing gradients. ``scaleRho`` must be tweaked to ensure proper gradient flow. Usually
        monitoring it for first few epochs is good enough.

        Usage:

        .. code-block:: python
            
            net = Network(netParams)
            net.gradFlow(path_to_save)
        '''
        gradNorm = lambda x: torch.linalg.norm(x).item()/torch.numel(x)
        grad = []

        for l in self.blocks:
            # print(l)
            if l.gradLog is True:
                if l.weightNorm is True:
                    grad.append(gradNorm(l.weightOp.weight_g.grad))
                else:
                    grad.append(gradNorm(l.weightOp.weight.grad))

        plt.figure()
        plt.semilogy(grad)
        plt.savefig(path + 'gradFlow.png')
        plt.close()
//  [python]
def test_gradient_clipping_by_norm(tmp_path, precision):
    """Test gradient clipping by norm."""
    trainer = Trainer(
        default_root_dir=tmp_path,
        max_steps=1,
        max_epochs=1,
        accelerator="auto",
        devices=1,
        precision=precision,
        gradient_clip_algorithm="norm",
        gradient_clip_val=0.05,
    )

    class TestModel(ClassificationModel):
        def configure_gradient_clipping(self, *args, **kwargs):
            super().configure_gradient_clipping(*args, **kwargs)
            # test that gradient is clipped correctly
            parameters = self.parameters()
            grad_norm = torch.norm(torch.stack([torch.norm(p.grad.detach(), 2) for p in parameters]), 2)
            torch.testing.assert_close(grad_norm, torch.tensor(0.05, device=self.device))
            self.assertion_called = True

    model = TestModel()
    trainer.fit(model, ClassifDataModule())
    assert model.assertion_called
//  [python]
def test_gradient_clipping_by_norm(tmp_path, precision):
    """Test gradient clipping by norm."""
    trainer = Trainer(
        default_root_dir=tmp_path,
        max_steps=1,
        max_epochs=1,
        accelerator="auto",
        devices=1,
        precision=precision,
        gradient_clip_algorithm="norm",
        gradient_clip_val=0.05,
    )

    class TestModel(ClassificationModel):
        def configure_gradient_clipping(self, *args, **kwargs):
            super().configure_gradient_clipping(*args, **kwargs)
            # test that gradient is clipped correctly
            parameters = self.parameters()
            grad_norm = torch.linalg.norm(torch.stack([torch.linalg.norm(p.grad.detach(), 2) for p in parameters]), 2)
            torch.testing.assert_close(grad_norm, torch.tensor(0.05, device=self.device))
            self.assertion_called = True

    model = TestModel()
    trainer.fit(model, ClassifDataModule())
    assert model.assertion_called
//  [python]
def plot_edge_direction(mesh_pos, edges):

    senders = edges[:, 0]
    receivers = edges[:, 1]

    edge_vec = torch.index_select(mesh_pos, 0, senders) - torch.index_select(
        mesh_pos, 0, receivers
    )
    e_x = torch.cat(
        (torch.ones(edge_vec.shape[0], 1), (torch.zeros(edge_vec.shape[0], 1))), dim=1
    )
    e_y = torch.cat(
        (torch.zeros(edge_vec.shape[0], 1), (torch.ones(edge_vec.shape[0], 1))), dim=1
    )

    edge_vec_dot_ex = edge_vec[:, 0] * e_x[:, 0] + edge_vec[:, 1] * e_x[:, 1]
    edge_vec_dot_ey = edge_vec[:, 0] * e_y[:, 0] + edge_vec[:, 1] * e_y[:, 1]

    cosu = edge_vec_dot_ex / ((torch.norm(edge_vec, dim=1) * torch.norm(e_x, dim=1)))
    cosv = edge_vec_dot_ey / ((torch.norm(edge_vec, dim=1) * torch.norm(e_y, dim=1)))

    plt.quiver(
        torch.index_select(mesh_pos[:, 0:1], 0, senders),
        torch.index_select(mesh_pos[:, 1:2], 0, senders),
        edge_vec[:, 0],
        edge_vec[:, 1],
        units="height",
        scale=1.2,
        width=0.0025,
    )

    plt.show()
//  [python]
def plot_edge_direction(mesh_pos, edges):

    senders = edges[:, 0]
    receivers = edges[:, 1]

    edge_vec = torch.index_select(mesh_pos, 0, senders) - torch.index_select(
        mesh_pos, 0, receivers
    )
    e_x = torch.cat(
        (torch.ones(edge_vec.shape[0], 1), (torch.zeros(edge_vec.shape[0], 1))), dim=1
    )
    e_y = torch.cat(
        (torch.zeros(edge_vec.shape[0], 1), (torch.ones(edge_vec.shape[0], 1))), dim=1
    )

    edge_vec_dot_ex = edge_vec[:, 0] * e_x[:, 0] + edge_vec[:, 1] * e_x[:, 1]
    edge_vec_dot_ey = edge_vec[:, 0] * e_y[:, 0] + edge_vec[:, 1] * e_y[:, 1]

    cosu = edge_vec_dot_ex / ((torch.linalg.norm(edge_vec, dim=1) * torch.linalg.norm(e_x, dim=1)))
    cosv = edge_vec_dot_ey / ((torch.linalg.norm(edge_vec, dim=1) * torch.linalg.norm(e_y, dim=1)))

    plt.quiver(
        torch.index_select(mesh_pos[:, 0:1], 0, senders),
        torch.index_select(mesh_pos[:, 1:2], 0, senders),
        edge_vec[:, 0],
        edge_vec[:, 1],
        units="height",
        scale=1.2,
        width=0.0025,
    )

    plt.show()
//  [python]
def test_coordinate_identity_transform_rectangular(self):
        dims = (10, 5, 256, 256)
        x = torch.randn(dims, device=DEVICE)
        y = torch.randn(dims, device=DEVICE)

        z = po.tools.polar_to_rectangular(*po.tools.rectangular_to_polar(torch.complex(x, y)))

        assert torch.norm(x - z.real) < 1e-3
        assert torch.norm(y - z.imag) < 1e-3
//  [python]
def test_coordinate_identity_transform_rectangular(self):
        dims = (10, 5, 256, 256)
        x = torch.randn(dims, device=DEVICE)
        y = torch.randn(dims, device=DEVICE)

        z = po.tools.polar_to_rectangular(*po.tools.rectangular_to_polar(torch.complex(x, y)))

        assert torch.linalg.norm(x - z.real) < 1e-3
        assert torch.linalg.norm(y - z.imag) < 1e-3
//  [python]
def norm_calculator(model):
    model_norm = 0
    for param_index, param in enumerate(model.parameters()):
        model_norm += torch.norm(param) ** 2
    return torch.sqrt(model_norm).item()
//  [python]
def norm_calculator(model):
    model_norm = 0
    for param_index, param in enumerate(model.parameters()):
        model_norm += torch.linalg.norm(param) ** 2
    return torch.sqrt(model_norm).item()
//  [python]
def is_pose_plausible(poses, joint_info):
    """Check poses for very basic plausibility, mainly to deal with false positive detections
    that don't yield a meaningful pose prediction.
    """
    mean_bones = (
        ds3d.get_dataset(FLAGS.bone_length_dataset).train_bones
        if FLAGS.bone_length_dataset
        else spu.load_pickle(FLAGS.bone_length_file))
    joint2bone_mat = posepile.joint_info.get_joint2bone_mat(joint_info)
    bones = joint2bone_mat @ poses[..., :joint_info.n_joints, :]
    bone_lengths = torch.norm(bones, dim=-1)
    bone_length_relative = bone_lengths / mean_bones
    bone_length_diff = torch.abs(bone_lengths - mean_bones)

    relsmall = bone_length_relative < 0.1
    relbig = bone_length_relative > 3
    absdiffbig = bone_length_diff > 300
    is_implausible = torch.any(
        torch.logical_and(torch.logical_or(relbig, relsmall), absdiffbig), dim=-1)
    return torch.logical_not(is_implausible)
//  [python]
def is_pose_plausible(poses, joint_info):
    """Check poses for very basic plausibility, mainly to deal with false positive detections
    that don't yield a meaningful pose prediction.
    """
    mean_bones = (
        ds3d.get_dataset(FLAGS.bone_length_dataset).train_bones
        if FLAGS.bone_length_dataset
        else spu.load_pickle(FLAGS.bone_length_file))
    joint2bone_mat = posepile.joint_info.get_joint2bone_mat(joint_info)
    bones = joint2bone_mat @ poses[..., :joint_info.n_joints, :]
    bone_lengths = torch.linalg.norm(bones, dim=-1)
    bone_length_relative = bone_lengths / mean_bones
    bone_length_diff = torch.abs(bone_lengths - mean_bones)

    relsmall = bone_length_relative < 0.1
    relbig = bone_length_relative > 3
    absdiffbig = bone_length_diff > 300
    is_implausible = torch.any(
        torch.logical_and(torch.logical_or(relbig, relsmall), absdiffbig), dim=-1)
    return torch.logical_not(is_implausible)
//  [python]
def test_coordinate_identity_transform_polar(self):
        dims = (10, 5, 256, 256)

        # ensure vec len a is non-zero by adding .1 and then re-normalizing
        a = torch.rand(dims, device=DEVICE) + 0.1
        a = a / a.max()
        b = po.tools.rescale(torch.randn(dims, device=DEVICE), -pi / 2, pi / 2)

        A, B = po.tools.rectangular_to_polar(po.tools.polar_to_rectangular(a, b))

        assert torch.norm(a - A) < 1e-3
        assert torch.norm(b - B) < 1e-3
//  [python]
def test_coordinate_identity_transform_polar(self):
        dims = (10, 5, 256, 256)

        # ensure vec len a is non-zero by adding .1 and then re-normalizing
        a = torch.rand(dims, device=DEVICE) + 0.1
        a = a / a.max()
        b = po.tools.rescale(torch.randn(dims, device=DEVICE), -pi / 2, pi / 2)

        A, B = po.tools.rectangular_to_polar(po.tools.polar_to_rectangular(a, b))

        assert torch.linalg.norm(a - A) < 1e-3
        assert torch.linalg.norm(b - B) < 1e-3
//  [python]
def ndc2dist(ndc_pts, cos_angle):
    dists = torch.norm(ndc_pts[:, 1:] - ndc_pts[:, :-1], dim=-1)
    dists = torch.cat([dists, 1e10 * cos_angle.unsqueeze(-1)], -1)  # [N_rays, N_samples]
    return dists
//  [python]
def ndc2dist(ndc_pts, cos_angle):
    dists = torch.linalg.norm(ndc_pts[:, 1:] - ndc_pts[:, :-1], dim=-1)
    dists = torch.cat([dists, 1e10 * cos_angle.unsqueeze(-1)], -1)  # [N_rays, N_samples]
    return dists
//  [python]
def _sample_orthonormal_to(self, mu, dim):
        """Sample point on sphere orthogonal to mu.
        """
        v = Variable(torch.randn(dim)).cuda()
        rescale_value = mu.dot(v) / mu.norm()
        proj_mu_v = mu * rescale_value.expand(dim)
        ortho = v - proj_mu_v
        ortho_norm = torch.norm(ortho)
        return ortho / ortho_norm.expand_as(ortho)
//  [python]
def _sample_orthonormal_to(self, mu, dim):
        """Sample point on sphere orthogonal to mu.
        """
        v = Variable(torch.randn(dim)).cuda()
        rescale_value = mu.dot(v) / torch.linalg.norm(mu)
        proj_mu_v = mu * rescale_value.expand(dim)
        ortho = v - proj_mu_v
        ortho_norm = torch.linalg.norm(ortho)
        return ortho / ortho_norm.expand_as(ortho)
//  [python]
def bislerp(samples, width, height):
    def slerp(b1, b2, r):
        '''slerps batches b1, b2 according to ratio r, batches should be flat e.g. NxC'''
        
        c = b1.shape[-1]

        #norms
        b1_norms = torch.norm(b1, dim=-1, keepdim=True)
        b2_norms = torch.norm(b2, dim=-1, keepdim=True)

        #normalize
        b1_normalized = b1 / b1_norms
        b2_normalized = b2 / b2_norms

        #zero when norms are zero
        b1_normalized[b1_norms.expand(-1,c) == 0.0] = 0.0
        b2_normalized[b2_norms.expand(-1,c) == 0.0] = 0.0

        #slerp
        dot = (b1_normalized*b2_normalized).sum(1)
        omega = torch.acos(dot)
        so = torch.sin(omega)

        #technically not mathematically correct, but more pleasing?
        res = (torch.sin((1.0-r.squeeze(1))*omega)/so).unsqueeze(1)*b1_normalized + (torch.sin(r.squeeze(1)*omega)/so).unsqueeze(1) * b2_normalized
        res *= (b1_norms * (1.0-r) + b2_norms * r).expand(-1,c)

        #edge cases for same or polar opposites
        res[dot > 1 - 1e-5] = b1[dot > 1 - 1e-5] 
        res[dot < 1e-5 - 1] = (b1 * (1.0-r) + b2 * r)[dot < 1e-5 - 1]
        return res
    
    def generate_bilinear_data(length_old, length_new, device):
        coords_1 = torch.arange(length_old, dtype=torch.float32, device=device).reshape((1,1,1,-1))
        coords_1 = torch.nn.functional.interpolate(coords_1, size=(1, length_new), mode="bilinear")
        ratios = coords_1 - coords_1.floor()
        coords_1 = coords_1.to(torch.int64)
        
        coords_2 = torch.arange(length_old, dtype=torch.float32, device=device).reshape((1,1,1,-1)) + 1
        coords_2[:,:,:,-1] -= 1
        coords_2 = torch.nn.functional.interpolate(coords_2, size=(1, length_new), mode="bilinear")
        coords_2 = coords_2.to(torch.int64)
        return ratios, coords_1, coords_2

    orig_dtype = samples.dtype
    samples = samples.float()
    n,c,h,w = samples.shape
    h_new, w_new = (height, width)
    
    #linear w
    ratios, coords_1, coords_2 = generate_bilinear_data(w, w_new, samples.device)
    coords_1 = coords_1.expand((n, c, h, -1))
    coords_2 = coords_2.expand((n, c, h, -1))
    ratios = ratios.expand((n, 1, h, -1))

    pass_1 = samples.gather(-1,coords_1).movedim(1, -1).reshape((-1,c))
    pass_2 = samples.gather(-1,coords_2).movedim(1, -1).reshape((-1,c))
    ratios = ratios.movedim(1, -1).reshape((-1,1))

    result = slerp(pass_1, pass_2, ratios)
    result = result.reshape(n, h, w_new, c).movedim(-1, 1)

    #linear h
    ratios, coords_1, coords_2 = generate_bilinear_data(h, h_new, samples.device)
    coords_1 = coords_1.reshape((1,1,-1,1)).expand((n, c, -1, w_new))
    coords_2 = coords_2.reshape((1,1,-1,1)).expand((n, c, -1, w_new))
    ratios = ratios.reshape((1,1,-1,1)).expand((n, 1, -1, w_new))

    pass_1 = result.gather(-2,coords_1).movedim(1, -1).reshape((-1,c))
    pass_2 = result.gather(-2,coords_2).movedim(1, -1).reshape((-1,c))
    ratios = ratios.movedim(1, -1).reshape((-1,1))

    result = slerp(pass_1, pass_2, ratios)
    result = result.reshape(n, h_new, w_new, c).movedim(-1, 1)
    return result.to(orig_dtype)
//  [python]
def bislerp(samples, width, height):
    def slerp(b1, b2, r):
        '''slerps batches b1, b2 according to ratio r, batches should be flat e.g. NxC'''
        
        c = b1.shape[-1]

        #norms
        b1_norms = torch.linalg.norm(b1, dim=-1, keepdim=True)
        b2_norms = torch.linalg.norm(b2, dim=-1, keepdim=True)

        #normalize
        b1_normalized = b1 / b1_norms
        b2_normalized = b2 / b2_norms

        #zero when norms are zero
        b1_normalized[b1_norms.expand(-1,c) == 0.0] = 0.0
        b2_normalized[b2_norms.expand(-1,c) == 0.0] = 0.0

        #slerp
        dot = (b1_normalized*b2_normalized).sum(1)
        omega = torch.acos(dot)
        so = torch.sin(omega)

        #technically not mathematically correct, but more pleasing?
        res = (torch.sin((1.0-r.squeeze(1))*omega)/so).unsqueeze(1)*b1_normalized + (torch.sin(r.squeeze(1)*omega)/so).unsqueeze(1) * b2_normalized
        res *= (b1_norms * (1.0-r) + b2_norms * r).expand(-1,c)

        #edge cases for same or polar opposites
        res[dot > 1 - 1e-5] = b1[dot > 1 - 1e-5] 
        res[dot < 1e-5 - 1] = (b1 * (1.0-r) + b2 * r)[dot < 1e-5 - 1]
        return res
    
    def generate_bilinear_data(length_old, length_new, device):
        coords_1 = torch.arange(length_old, dtype=torch.float32, device=device).reshape((1,1,1,-1))
        coords_1 = torch.nn.functional.interpolate(coords_1, size=(1, length_new), mode="bilinear")
        ratios = coords_1 - coords_1.floor()
        coords_1 = coords_1.to(torch.int64)
        
        coords_2 = torch.arange(length_old, dtype=torch.float32, device=device).reshape((1,1,1,-1)) + 1
        coords_2[:,:,:,-1] -= 1
        coords_2 = torch.nn.functional.interpolate(coords_2, size=(1, length_new), mode="bilinear")
        coords_2 = coords_2.to(torch.int64)
        return ratios, coords_1, coords_2

    orig_dtype = samples.dtype
    samples = samples.float()
    n,c,h,w = samples.shape
    h_new, w_new = (height, width)
    
    #linear w
    ratios, coords_1, coords_2 = generate_bilinear_data(w, w_new, samples.device)
    coords_1 = coords_1.expand((n, c, h, -1))
    coords_2 = coords_2.expand((n, c, h, -1))
    ratios = ratios.expand((n, 1, h, -1))

    pass_1 = samples.gather(-1,coords_1).movedim(1, -1).reshape((-1,c))
    pass_2 = samples.gather(-1,coords_2).movedim(1, -1).reshape((-1,c))
    ratios = ratios.movedim(1, -1).reshape((-1,1))

    result = slerp(pass_1, pass_2, ratios)
    result = result.reshape(n, h, w_new, c).movedim(-1, 1)

    #linear h
    ratios, coords_1, coords_2 = generate_bilinear_data(h, h_new, samples.device)
    coords_1 = coords_1.reshape((1,1,-1,1)).expand((n, c, -1, w_new))
    coords_2 = coords_2.reshape((1,1,-1,1)).expand((n, c, -1, w_new))
    ratios = ratios.reshape((1,1,-1,1)).expand((n, 1, -1, w_new))

    pass_1 = result.gather(-2,coords_1).movedim(1, -1).reshape((-1,c))
    pass_2 = result.gather(-2,coords_2).movedim(1, -1).reshape((-1,c))
    ratios = ratios.movedim(1, -1).reshape((-1,1))

    result = slerp(pass_1, pass_2, ratios)
    result = result.reshape(n, h_new, w_new, c).movedim(-1, 1)
    return result.to(orig_dtype)
//  [python]
def init_meta(self):
        if self.solved_meta is not None:
            return self.solved_meta
        skill = torch.randn(self.skill_dim)
        skill = skill / torch.norm(skill)
        skill = skill.cpu().numpy()
        meta = OrderedDict()
        meta['skill'] = skill
        return meta
//  [python]
def init_meta(self):
        if self.solved_meta is not None:
            return self.solved_meta
        skill = torch.randn(self.skill_dim)
        skill = skill / torch.linalg.norm(skill)
        skill = skill.cpu().numpy()
        meta = OrderedDict()
        meta['skill'] = skill
        return meta
//  [python]
def get_kmeans_mu(
        self, x, n_centers, init_times: int = 50, min_delta: float = 1e-3
    ) -> Tensor:
        """Find an initial value for the mean. 
        
        Requires a threshold min_delta for the k-means algorithm to stop iterating.
        The algorithm is repeated init_times often, after which the best centerpoint is returned.

        Args:
            x (Tensor): A tensor of shape of (n, d) or (n, 1, d).
            init_times (int): Number of times to run the k-means algorithm.
            min_delta (float): Minimum change in the loss function to stop the k-means algorithm.
        
        Returns:
            Tensor: A tensor of shape (1, k, d).
        """
        if len(x.size()) == 3:
            x = x.squeeze(1)
        x_min, x_max = x.min(), x.max()
        x = (x - x_min) / (x_max - x_min)

        min_cost = np.inf

        center = x[
            np.random.choice(np.arange(x.shape[0]), size=n_centers, replace=False), ...,
        ]

        for i in range(init_times):
            tmp_center = x[
                np.random.choice(np.arange(x.shape[0]), size=n_centers, replace=False),
                ...,
            ]
            l2_dis = torch.norm(
                (x.unsqueeze(1).repeat(1, n_centers, 1) - tmp_center), p=2, dim=2
            )
            l2_cls = torch.argmin(l2_dis, dim=1)

            cost = 0
            for c in range(n_centers):
                cost += torch.norm(x[l2_cls == c] - tmp_center[c], p=2, dim=1).mean()

            if cost < min_cost:
                min_cost = cost
                center = tmp_center

        delta = np.inf

        while delta > min_delta:
            l2_dis = torch.norm(
                (x.unsqueeze(1).repeat(1, n_centers, 1) - center), p=2, dim=2
            )
            l2_cls = torch.argmin(l2_dis, dim=1)
            center_old = center.clone()

            for c in range(n_centers):
                center[c] = x[l2_cls == c].mean(dim=0)

            delta = torch.norm((center_old - center), dim=1).max()

        return center.unsqueeze(0) * (x_max - x_min) + x_min
//  [python]
def get_kmeans_mu(
        self, x, n_centers, init_times: int = 50, min_delta: float = 1e-3
    ) -> Tensor:
        """Find an initial value for the mean. 
        
        Requires a threshold min_delta for the k-means algorithm to stop iterating.
        The algorithm is repeated init_times often, after which the best centerpoint is returned.

        Args:
            x (Tensor): A tensor of shape of (n, d) or (n, 1, d).
            init_times (int): Number of times to run the k-means algorithm.
            min_delta (float): Minimum change in the loss function to stop the k-means algorithm.
        
        Returns:
            Tensor: A tensor of shape (1, k, d).
        """
        if len(x.size()) == 3:
            x = x.squeeze(1)
        x_min, x_max = x.min(), x.max()
        x = (x - x_min) / (x_max - x_min)

        min_cost = np.inf

        center = x[
            np.random.choice(np.arange(x.shape[0]), size=n_centers, replace=False), ...,
        ]

        for i in range(init_times):
            tmp_center = x[
                np.random.choice(np.arange(x.shape[0]), size=n_centers, replace=False),
                ...,
            ]
            l2_dis = torch.linalg.norm(
                (x.unsqueeze(1).repeat(1, n_centers, 1) - tmp_center), ord=2, dim=2
            )
            l2_cls = torch.argmin(l2_dis, dim=1)

            cost = 0
            for c in range(n_centers):
                cost += torch.linalg.norm(x[l2_cls == c] - tmp_center[c], ord=2, dim=1).mean()

            if cost < min_cost:
                min_cost = cost
                center = tmp_center

        delta = np.inf

        while delta > min_delta:
            l2_dis = torch.linalg.norm(
                (x.unsqueeze(1).repeat(1, n_centers, 1) - center), ord=2, dim=2
            )
            l2_cls = torch.argmin(l2_dis, dim=1)
            center_old = center.clone()

            for c in range(n_centers):
                center[c] = x[l2_cls == c].mean(dim=0)

            delta = torch.linalg.norm((center_old - center), ord=2, dim=1).max()

        return center.unsqueeze(0) * (x_max - x_min) + x_min
//  [python]
def vmf_unif_sampler(self, mu):
        batch_size, id_dim = mu.size()
        result_list = []
        for i in range(batch_size):
            munorm = mu[i].norm().expand(id_dim)
            munoise = self.add_norm_noise(munorm, self.norm_eps)
            if float(mu[i].norm().data.cpu().numpy()) > 1e-10:
                # sample offset from center (on sphere) with spread kappa
                w = self._sample_weight(self.kappa, id_dim)
                wtorch = torch.autograd.Variable(w * torch.ones(id_dim)).cuda()

                # sample a point v on the unit sphere that's orthogonal to mu
                v = self._sample_orthonormal_to(mu[i] / munorm, id_dim)

                # compute new point
                scale_factr = torch.sqrt(torch.autograd.Variable(torch.ones(id_dim)).cuda() - torch.pow(wtorch, 2))
                orth_term = v * scale_factr
                muscale = mu[i] * wtorch / munorm
                sampled_vec = (orth_term + muscale) * munoise
            else:
                rand_draw = torch.autograd.Variable(torch.randn(id_dim))
                rand_draw = rand_draw / torch.norm(rand_draw, p=2).expand(id_dim)
                rand_norms = (torch.rand(1) * self.norm_eps).expand(id_dim)
                sampled_vec = rand_draw * torch.autograd.Variable(rand_norms).cuda()  # mu[i]
            result_list.append(sampled_vec)

        return torch.stack(result_list, 0)
//  [python]
def vmf_unif_sampler(self, mu):
        batch_size, id_dim = mu.size()
        result_list = []
        for i in range(batch_size):
            munorm = torch.linalg.norm(mu[i]).expand(id_dim)
            munoise = self.add_norm_noise(munorm, self.norm_eps)
            if float(torch.linalg.norm(mu[i]).data.cpu().numpy()) > 1e-10:
                # sample offset from center (on sphere) with spread kappa
                w = self._sample_weight(self.kappa, id_dim)
                wtorch = torch.autograd.Variable(w * torch.ones(id_dim)).cuda()

                # sample a point v on the unit sphere that's orthogonal to mu
                v = self._sample_orthonormal_to(mu[i] / munorm, id_dim)

                # compute new point
                scale_factr = torch.sqrt(torch.autograd.Variable(torch.ones(id_dim)).cuda() - torch.pow(wtorch, 2))
                orth_term = v * scale_factr
                muscale = mu[i] * wtorch / munorm
                sampled_vec = (orth_term + muscale) * munoise
            else:
                rand_draw = torch.autograd.Variable(torch.randn(id_dim))
                rand_draw = rand_draw / torch.linalg.norm(rand_draw, p=2).expand(id_dim)
                rand_norms = (torch.rand(1) * self.norm_eps).expand(id_dim)
                sampled_vec = rand_draw * torch.autograd.Variable(rand_norms).cuda()  # mu[i]
            result_list.append(sampled_vec)

        return torch.stack(result_list, 0)
//  [python]
def _calc_coverage(self, dset, a_tch, b_tch, rho=1):
        """
        This function calculates coverage.

        Args:
            dset:
                Dataset (train or test)
            a_tch:
                A torch
            b_tch:
                b torch

        Returns:
            Coverage
        """
        coverage = 0
        for datind in range(dset.shape[0]):
            coverage += torch.where(
                torch.norm((a_tch.T@torch.linalg.inv(a_tch@a_tch.T)) @ (dset[datind]-
                           b_tch))
                <= rho,
                1,
                0,
            )
        return coverage/dset.shape[0]
//  [python]
def _calc_coverage(self, dset, a_tch, b_tch, rho=1):
        """
        This function calculates coverage.

        Args:
            dset:
                Dataset (train or test)
            a_tch:
                A torch
            b_tch:
                b torch

        Returns:
            Coverage
        """
        coverage = 0
        for datind in range(dset.shape[0]):
            coverage += torch.where(
                torch.linalg.norm((a_tch.T@torch.linalg.inv(a_tch@a_tch.T)) @ (dset[datind]-
                           b_tch))
                <= rho,
                1,
                0,
            )
        return coverage/dset.shape[0]
//  [python]
def vmf_sampler(self, mu):
        mu = mu.cpu()
        batch_size, id_dim = mu.size()
        result_list = []
        for i in range(batch_size):
            munorm = mu[i].norm().expand(id_dim)  # TODO norm p=?
            if float(mu[i].norm().data.cpu().numpy()) > 1e-10:
                # sample offset from center (on sphere) with spread kappa
                # w = self._sample_weight(self.kappa, id_dim)     # TODO mine?

                w = vMF.sample_vmf_w(self.kappa, id_dim)

                wtorch = torch.autograd.Variable(w * torch.ones(id_dim))

                # sample a point v on the unit sphere that's orthogonal to mu
                v = self._sample_orthonormal_to(mu[i] / munorm, id_dim)
                # v= vMF.sample_vmf_v(mu[i])

                # compute new point
                scale_factr = torch.sqrt(Variable(torch.ones(id_dim)) - torch.pow(wtorch, 2))
                orth_term = v * scale_factr

                muscale = mu[i] * wtorch / munorm
                sampled_vec = (orth_term + muscale) * munorm
            else:
                rand_draw = Variable(torch.randn(id_dim))
                rand_draw = rand_draw / torch.norm(rand_draw, p=2).expand(id_dim)
                rand_norms = (torch.rand(1) * self.norm_eps).expand(id_dim)
                sampled_vec = rand_draw * Variable(rand_norms)  # mu[i]
            result_list.append(sampled_vec)

        return torch.stack(result_list, 0).cuda()
//  [python]
def vmf_sampler(self, mu):
        mu = mu.cpu()
        batch_size, id_dim = mu.size()
        result_list = []
        for i in range(batch_size):
            munorm = torch.linalg.norm(mu[i]).expand(id_dim)  # TODO norm p=?
            if float(torch.linalg.norm(mu[i]).data.cpu().numpy()) > 1e-10:
                # sample offset from center (on sphere) with spread kappa
                # w = self._sample_weight(self.kappa, id_dim)     # TODO mine?

                w = vMF.sample_vmf_w(self.kappa, id_dim)

                wtorch = torch.autograd.Variable(w * torch.ones(id_dim))

                # sample a point v on the unit sphere that's orthogonal to mu
                v = self._sample_orthonormal_to(mu[i] / munorm, id_dim)
                # v= vMF.sample_vmf_v(mu[i])

                # compute new point
                scale_factr = torch.sqrt(Variable(torch.ones(id_dim)) - torch.pow(wtorch, 2))
                orth_term = v * scale_factr

                muscale = mu[i] * wtorch / munorm
                sampled_vec = (orth_term + muscale) * munorm
            else:
                rand_draw = Variable(torch.randn(id_dim))
                rand_draw = rand_draw / torch.linalg.norm(rand_draw, p=2).expand(id_dim)
                rand_norms = (torch.rand(1) * self.norm_eps).expand(id_dim)
                sampled_vec = rand_draw * Variable(rand_norms)  # mu[i]
            result_list.append(sampled_vec)

        return torch.stack(result_list, 0).cuda()
//  [python]
def compute_intr_reward(self, seq, step):
        skill, next_obs = seq['skill'], seq['feat']
        
        B, T, _ = skill.shape
        skill = skill.reshape(B*T, -1)
        next_obs = next_obs.reshape(B*T, -1)

        # maxent reward
        with torch.no_grad():
            rep = self.aps(next_obs, norm=False)
            reward = self.aps_pbe(rep, cdist=True)
            intr_ent_reward = reward.reshape(-1, 1)

        # successor feature reward
        rep = rep / torch.norm(rep, dim=1, keepdim=True)
        intr_sf_reward = torch.einsum("bi,bi->b", skill, rep).reshape(-1, 1)

        reward = intr_ent_reward + intr_sf_reward

        return reward.reshape(B, T, 1).detach()
//  [python]
def compute_intr_reward(self, seq, step):
        skill, next_obs = seq['skill'], seq['feat']
        
        B, T, _ = skill.shape
        skill = skill.reshape(B*T, -1)
        next_obs = next_obs.reshape(B*T, -1)

        # maxent reward
        with torch.no_grad():
            rep = self.aps(next_obs, norm=False)
            reward = self.aps_pbe(rep, cdist=True)
            intr_ent_reward = reward.reshape(-1, 1)

        # successor feature reward
        rep = rep / torch.linalg.norm(rep, dim=1, keepdim=True)
        intr_sf_reward = torch.einsum("bi,bi->b", skill, rep).reshape(-1, 1)

        reward = intr_ent_reward + intr_sf_reward

        return reward.reshape(B, T, 1).detach()