jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def photo_level_tracking(Tji_init, aff_init, vals_i, Pi, dI_dT, img_j, intrinsics, photo_sigma, term_criteria):
  Tji = Tji_init.clone()
  aff = aff_init.clone()

  A_norm = 1.0/torch.as_tensor((img_j.shape[-1], img_j.shape[-2]), device=img_j.device, dtype=img_j.dtype)

  iter = 0
  done = False
  mean_err_prev = float("inf")
  while not done:
    Tji, aff, delta, mean_err, p_j, valid_mask, depth_j = tracking_iter(Tji, Pi, intrinsics, img_j, aff, vals_i, dI_dT, photo_sigma, A_norm)

    # print("Tracking: ", iter, mean_err.item())

    iter += 1
    delta_norm = torch.norm(delta)
    abs_decrease = mean_err_prev - mean_err
    rel_decrease = abs_decrease/mean_err_prev
    if iter >= term_criteria["max_iter"] \
      or abs_decrease < term_criteria["abs_tol"] \
        or delta_norm < term_criteria["delta_norm"] \
        or rel_decrease < term_criteria["rel_tol"]:
      done = True
      # print(iter, abs_decrease, delta_norm, rel_decrease)
    mean_err_prev = mean_err

  test_coords_target = swap_coords_xy(p_j)
  coords_j = test_coords_target[0:1,valid_mask[0,:],:]
  depths_j = depth_j[0:1,valid_mask[0,:],:]

  return Tji, aff, coords_j, depths_j
//  [python]
def photo_level_tracking(Tji_init, aff_init, vals_i, Pi, dI_dT, img_j, intrinsics, photo_sigma, term_criteria):
  Tji = Tji_init.clone()
  aff = aff_init.clone()

  A_norm = 1.0/torch.as_tensor((img_j.shape[-1], img_j.shape[-2]), device=img_j.device, dtype=img_j.dtype)

  iter = 0
  done = False
  mean_err_prev = float("inf")
  while not done:
    Tji, aff, delta, mean_err, p_j, valid_mask, depth_j = tracking_iter(Tji, Pi, intrinsics, img_j, aff, vals_i, dI_dT, photo_sigma, A_norm)

    # print("Tracking: ", iter, mean_err.item())

    iter += 1
    delta_norm = torch.linalg.norm(delta)
    abs_decrease = mean_err_prev - mean_err
    rel_decrease = abs_decrease/mean_err_prev
    if iter >= term_criteria["max_iter"] \
      or abs_decrease < term_criteria["abs_tol"] \
        or delta_norm < term_criteria["delta_norm"] \
        or rel_decrease < term_criteria["rel_tol"]:
      done = True
      # print(iter, abs_decrease, delta_norm, rel_decrease)
    mean_err_prev = mean_err

  test_coords_target = swap_coords_xy(p_j)
  coords_j = test_coords_target[0:1,valid_mask[0,:],:]
  depths_j = depth_j[0:1,valid_mask[0,:],:]

  return Tji, aff, coords_j, depths_j
//  [python]
def run_on_batch(self, batch):
        """ Run algorithm on batch of images in eval mode

            @param batch: a dictionary with the following keys:
                            - xyz: a [N x 3 x H x W] torch.FloatTensor

            @return fg_mask: a [N x H x W] torch.LongTensor with values in {0, 1, 2}
                    center_direction_prediction: a [N x 2 x H x W] torch.FloatTensor
                    object_centers: a list of [2 x num_objects] torch.IntTensor. This list has length N
                    initial_maks: a [N x H x W] torch.IntTensor
        """

        self.eval_mode()
        self.send_batch_to_device(batch)

        with torch.no_grad():

            # Apply model
            fg_logits, center_direction_prediction = self.model(batch['xyz'])

            # Foreground
            fg_probs = F.softmax(fg_logits, dim=1) # Shape: [N x 3 x H x W]
            fg_mask = torch.argmax(fg_probs, dim=1) # Shape: [N x H x W]

            # Center direction
            center_direction_prediction = center_direction_prediction / torch.norm(center_direction_prediction, 
                                                                                   dim=1, 
                                                                                   keepdim=True
                                                                                  ).clamp(min=1e-10)

            initial_masks, num_objects, object_centers_padded = \
                self.hough_voting_layer((fg_mask == OBJECTS_LABEL).int(), 
                                        center_direction_prediction)

        # Compute list of object centers
        width = initial_masks.shape[2]
        object_centers = []
        for i in range(initial_masks.shape[0]):
            object_centers.append(object_centers_padded[i, :, :num_objects[i]])

        return fg_mask, center_direction_prediction, object_centers, initial_masks
//  [python]
def run_on_batch(self, batch):
        """ Run algorithm on batch of images in eval mode

            @param batch: a dictionary with the following keys:
                            - xyz: a [N x 3 x H x W] torch.FloatTensor

            @return fg_mask: a [N x H x W] torch.LongTensor with values in {0, 1, 2}
                    center_direction_prediction: a [N x 2 x H x W] torch.FloatTensor
                    object_centers: a list of [2 x num_objects] torch.IntTensor. This list has length N
                    initial_maks: a [N x H x W] torch.IntTensor
        """

        self.eval_mode()
        self.send_batch_to_device(batch)

        with torch.no_grad():

            # Apply model
            fg_logits, center_direction_prediction = self.model(batch['xyz'])

            # Foreground
            fg_probs = F.softmax(fg_logits, dim=1) # Shape: [N x 3 x H x W]
            fg_mask = torch.argmax(fg_probs, dim=1) # Shape: [N x H x W]

            # Center direction
            center_direction_prediction = center_direction_prediction / torch.linalg.norm(center_direction_prediction, 
                                                                                   dim=1, 
                                                                                   keepdim=True
                                                                                  ).clamp(min=1e-10)

            initial_masks, num_objects, object_centers_padded = \
                self.hough_voting_layer((fg_mask == OBJECTS_LABEL).int(), 
                                        center_direction_prediction)

        # Compute list of object centers
        width = initial_masks.shape[2]
        object_centers = []
        for i in range(initial_masks.shape[0]):
            object_centers.append(object_centers_padded[i, :, :num_objects[i]])

        return fg_mask, center_direction_prediction, object_centers, initial_masks
//  [python]
def calc_error_rate_i(pred_landmark_coords, gt_landmark_coords):
    error_single = torch.norm(pred_landmark_coords - gt_landmark_coords, dim=1)
    error = torch.mean(error_single)
    return error, error_single
//  [python]
def calc_error_rate_i(pred_landmark_coords, gt_landmark_coords):
    error_single = torch.linalg.norm(pred_landmark_coords - gt_landmark_coords, dim=1)
    error = torch.mean(error_single)
    return error, error_single
//  [python]
def compare_impl():
    """
    lstsq : error 2.23668351395645e-07 erank 4.248441219329834 discrepancy 4760.62841796875
    scipy : error 4.3144709138687176e-07 erank 5.015336036682129 discrepancy 0.38698798418045044
    svd : error 2.6252257612213725e-07 erank 5.567287921905518 discrepancy 11.784327507019043
    spectral : error 3.1967988434189465e-07 erank 4.978618144989014 discrepancy 0.0020115238148719072
    """

    bad_sigmas = torch.load('test/bad_sigmas.pt')
    H = bad_sigmas['H']
    sigma = bad_sigmas['sigma']

    # X = u.lyapunov_spectral(H, sigma)

    def print_stats(tag, X):
        error = torch.norm(H @ X + X @ H - sigma)
        discrepancy = torch.max(abs(X - X.t()) / X)
        print(tag, ": error", error.item(), "erank", u.erank(X).item(), 'discrepancy', discrepancy.item())

    H0, sigma0 = u.to_numpys(H, sigma)
    X0 = scipy.linalg.solve_lyapunov(H0, sigma0)

    print_stats('lstsq', u.lyapunov_lstsq(H, sigma))
    print_stats('scipy', torch.tensor(X0))
    print_stats('svd', u.lyapunov_svd(H, sigma))
    print_stats('spectral', u.lyapunov_spectral(H, sigma))
//  [python]
def compare_impl():
    """
    lstsq : error 2.23668351395645e-07 erank 4.248441219329834 discrepancy 4760.62841796875
    scipy : error 4.3144709138687176e-07 erank 5.015336036682129 discrepancy 0.38698798418045044
    svd : error 2.6252257612213725e-07 erank 5.567287921905518 discrepancy 11.784327507019043
    spectral : error 3.1967988434189465e-07 erank 4.978618144989014 discrepancy 0.0020115238148719072
    """

    bad_sigmas = torch.load('test/bad_sigmas.pt')
    H = bad_sigmas['H']
    sigma = bad_sigmas['sigma']

    # X = u.lyapunov_spectral(H, sigma)

    def print_stats(tag, X):
        error = torch.linalg.norm(H @ X + X @ H - sigma)
        discrepancy = torch.max(abs(X - X.t()) / X)
        print(tag, ": error", error.item(), "erank", u.erank(X).item(), 'discrepancy', discrepancy.item())

    H0, sigma0 = u.to_numpys(H, sigma)
    X0 = scipy.linalg.solve_lyapunov(H0, sigma0)

    print_stats('lstsq', u.lyapunov_lstsq(H, sigma))
    print_stats('scipy', torch.tensor(X0))
    print_stats('svd', u.lyapunov_svd(H, sigma))
    print_stats('spectral', u.lyapunov_spectral(H, sigma))
//  [python]
def get_3D_FAN_NME(pred_landmark, gt_landmark):
    NME_single = torch.sum(torch.norm(pred_landmark - gt_landmark, dim=2), 0)
    NME = torch.mean(NME_single)
    return NME, NME_single
//  [python]
def get_3D_FAN_NME(pred_landmark, gt_landmark):
    NME_single = torch.sum(torch.linalg.norm(pred_landmark - gt_landmark, dim=2), 0)
    NME = torch.mean(NME_single)
    return NME, NME_single
//  [python]
def test_ema_fp32(self):
        model = DummyModule().cuda().half()
        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).cuda()
        y = model(x.half())
        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()).half().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()).half(),
            )
//  [python]
def test_ema_fp32(self):
        model = DummyModule().cuda().half()
        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).cuda()
        y = model(x.half())
        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()).half().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()).half(),
            )
//  [python]
def test_ema_fp16(self):
        model = DummyModule().cuda().half()
        optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
        state = deepcopy(model.state_dict())
        config = EMAConfig(ema_fp32=False)
        ema = EMA(model, config)

        # Since fp32 params is not used, it should be of size 0
        self.assertEqual(len(ema.fp32_params), 0)

        x = torch.randn(32).cuda()
        y = model(x.half())
        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

            # EMA update is done in fp16, and hence the EMA param must be
            # closer to the EMA update done in fp16 than in fp32.
            self.assertLessEqual(
                torch.norm(
                    ema_param.float() -
                    (config.ema_decay * prev_param + (1 - config.ema_decay) * param).float()
                ),
                torch.norm(
                    ema_param.float() -
                    (config.ema_decay * prev_param.float() + (1 - config.ema_decay) * param.float()).half().float()
                ),
            )
            self.assertTorchAllClose(
                ema_param,
                config.ema_decay * prev_param + (1 - config.ema_decay) * param,
            )

        # Since fp32 params is not used, it should be of size 0
        self.assertEqual(len(ema.fp32_params), 0)
//  [python]
def test_ema_fp16(self):
        model = DummyModule().cuda().half()
        optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
        state = deepcopy(model.state_dict())
        config = EMAConfig(ema_fp32=False)
        ema = EMA(model, config)

        # Since fp32 params is not used, it should be of size 0
        self.assertEqual(len(ema.fp32_params), 0)

        x = torch.randn(32).cuda()
        y = model(x.half())
        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

            # EMA update is done in fp16, and hence the EMA param must be
            # closer to the EMA update done in fp16 than in fp32.
            self.assertLessEqual(
                torch.linalg.norm(
                    ema_param.float() -
                    (config.ema_decay * prev_param + (1 - config.ema_decay) * param).float()
                ),
                torch.linalg.norm(
                    ema_param.float() -
                    (config.ema_decay * prev_param.float() + (1 - config.ema_decay) * param.float()).half().float()
                ),
            )
            self.assertTorchAllClose(
                ema_param,
                config.ema_decay * prev_param + (1 - config.ema_decay) * param,
            )

        # Since fp32 params is not used, it should be of size 0
        self.assertEqual(len(ema.fp32_params), 0)
//  [python]
def get_novel_view_around_bbox(opt, N=120, scale=1):
    # create surrounding view point around bbox
    bbox = torch.tensor(opt.data.scene_bbox).view(2,3).to(opt.device)
    center = torch.mean(bbox, dim=0)
    norm = torch.norm(bbox[0]-bbox[1])
    radius = norm * 0.2 * scale
    theta = torch.arange(N)/N*2*np.pi
    theta = theta.to(opt.device)
    camera_dist = norm * 0.6 * scale

    bbox_height = torch.abs(bbox[1,1] - bbox[0,1])
    h_above_center = bbox_height * 0.1 * scale
    pitch_angle = torch.atan(h_above_center / camera_dist )
    R_x = angle_to_rotation_matrix(pitch_angle ,"X")
    R_y = angle_to_rotation_matrix(theta, "Y")
    R_dwn_90 = angle_to_rotation_matrix(torch.full_like(theta, np.pi/2), "X")

    pose_x = torch.zeros_like(theta)
    pose_z = torch.full_like(theta, camera_dist)
    pose_y = torch.full_like(theta, h_above_center)

    novel_pose = pose(R=R_x@R_y@R_dwn_90, t=torch.stack([pose_x, pose_y, pose_z], dim=-1))
    return novel_pose
//  [python]
def get_novel_view_around_bbox(opt, N=120, scale=1):
    # create surrounding view point around bbox
    bbox = torch.tensor(opt.data.scene_bbox).view(2,3).to(opt.device)
    center = torch.mean(bbox, dim=0)
    norm = torch.linalg.norm(bbox[0]-bbox[1])
    radius = norm * 0.2 * scale
    theta = torch.arange(N)/N*2*np.pi
    theta = theta.to(opt.device)
    camera_dist = norm * 0.6 * scale

    bbox_height = torch.abs(bbox[1,1] - bbox[0,1])
    h_above_center = bbox_height * 0.1 * scale
    pitch_angle = torch.atan(h_above_center / camera_dist )
    R_x = angle_to_rotation_matrix(pitch_angle ,"X")
    R_y = angle_to_rotation_matrix(theta, "Y")
    R_dwn_90 = angle_to_rotation_matrix(torch.full_like(theta, np.pi/2), "X")

    pose_x = torch.zeros_like(theta)
    pose_z = torch.full_like(theta, camera_dist)
    pose_y = torch.full_like(theta, h_above_center)

    novel_pose = pose(R=R_x@R_y@R_dwn_90, t=torch.stack([pose_x, pose_y, pose_z], dim=-1))
    return novel_pose
//  [python]
def test_SYMEIG_random():
    m = 50
    M = torch.rand(m, m, dtype=torch.float64)
    M = 0.5 * (M + M.t())

    D, U = SYMEIG.apply(M)
    assert torch.norm(M - U @ torch.diag(D) @ U.t()) < D[0] * (m ** 2) * 1e-14

    # since we always assume matrix M to be symmetric, the finite difference
    # perturbations should be symmetric as well
    M.requires_grad_(True)

    def force_sym_eig(M):
        M = 0.5 * (M + M.t())
        return SYMEIG.apply(M)
    assert(torch.autograd.gradcheck(force_sym_eig, M, eps=1e-6, atol=1e-4))
//  [python]
def test_SYMEIG_random():
    m = 50
    M = torch.rand(m, m, dtype=torch.float64)
    M = 0.5 * (M + M.t())

    D, U = SYMEIG.apply(M)
    assert torch.linalg.norm(M - U @ torch.diag(D) @ U.t()) < D[0] * (m ** 2) * 1e-14

    # since we always assume matrix M to be symmetric, the finite difference
    # perturbations should be symmetric as well
    M.requires_grad_(True)

    def force_sym_eig(M):
        M = 0.5 * (M + M.t())
        return SYMEIG.apply(M)
    assert(torch.autograd.gradcheck(force_sym_eig, M, eps=1e-6, atol=1e-4))
//  [python]
def forward(self, mean, cov=None, *, decayscales, render_color=True):
    """Run MLP. mean is [D, *batch, 3] and cov is [D, *batch, 3, 3]."""
    batch_shape = mean.shape[:-1]
    num_points = np.prod(batch_shape)
    assert mean.shape[-1] == 3

    x = mean.reshape(num_points, 3)

    # Canonical network. If enabled, use IPE.
    assert cov is not None
    if not isinstance(cov, float):
      cov = cov.reshape(num_points, 3, 3)
    emb = self.emb(x, cov, decayscales["density"])
    features = self.layers_density(emb)
    d, features = features[Ellipsis, :1], features[Ellipsis, 1:]
    assert d.dtype == torch.float32

    # Parameterize canonical frame VolSDF distance as offset from sphere.
    if self.parameterization == "volsdf" and self.base_radius is not None:
      d = d + torch.norm(x, dim=-1, keepdim=True) - self.base_radius

    # Appearance network.
    if render_color:
      raw_rgb = self.layers_color(features)
      raw_output = torch.cat([raw_rgb, d], dim=-1)
      return raw_output.view(*batch_shape, 4)

    return d.view(*batch_shape)
//  [python]
def forward(self, mean, cov=None, *, decayscales, render_color=True):
    """Run MLP. mean is [D, *batch, 3] and cov is [D, *batch, 3, 3]."""
    batch_shape = mean.shape[:-1]
    num_points = np.prod(batch_shape)
    assert mean.shape[-1] == 3

    x = mean.reshape(num_points, 3)

    # Canonical network. If enabled, use IPE.
    assert cov is not None
    if not isinstance(cov, float):
      cov = cov.reshape(num_points, 3, 3)
    emb = self.emb(x, cov, decayscales["density"])
    features = self.layers_density(emb)
    d, features = features[Ellipsis, :1], features[Ellipsis, 1:]
    assert d.dtype == torch.float32

    # Parameterize canonical frame VolSDF distance as offset from sphere.
    if self.parameterization == "volsdf" and self.base_radius is not None:
      d = d + torch.linalg.norm(x, dim=-1, keepdim=True) - self.base_radius

    # Appearance network.
    if render_color:
      raw_rgb = self.layers_color(features)
      raw_output = torch.cat([raw_rgb, d], dim=-1)
      return raw_output.view(*batch_shape, 4)

    return d.view(*batch_shape)
//  [python]
def test_SYMEIG_3x3degenerate():
    M = torch.zeros((3, 3), dtype=torch.float64)
    M[0, 1] = M[0, 2] = M[1, 2] = 1.
    M = 0.5 * (M + M.t())
    print(M)

    D, U = SYMEIG.apply(M)
    assert torch.norm(M - U @ torch.diag(D) @ U.t()) < D[0] * (M.size()[0] ** 2) * 1e-14

    M.requires_grad_(True)
    torch.set_printoptions(precision=9)

    def force_sym_eig(M):
        M = 0.5 * (M + M.t())
        print(M)
        D, U = SYMEIG.apply(M)
        return U
    assert(torch.autograd.gradcheck(force_sym_eig, M, eps=1e-6, atol=1e-4))
//  [python]
def test_SYMEIG_3x3degenerate():
    M = torch.zeros((3, 3), dtype=torch.float64)
    M[0, 1] = M[0, 2] = M[1, 2] = 1.
    M = 0.5 * (M + M.t())
    print(M)

    D, U = SYMEIG.apply(M)
    assert torch.linalg.norm(M - U @ torch.diag(D) @ U.t()) < D[0] * (M.size()[0] ** 2) * 1e-14

    M.requires_grad_(True)
    torch.set_printoptions(precision=9)

    def force_sym_eig(M):
        M = 0.5 * (M + M.t())
        print(M)
        D, U = SYMEIG.apply(M)
        return U
    assert(torch.autograd.gradcheck(force_sym_eig, M, eps=1e-6, atol=1e-4))
//  [python]
def test_lyapunov():
    """Test that scipy lyapunov solver works correctly."""
    d = 2
    n = 3
    torch.set_default_dtype(torch.float32)

    model = Net(d)

    w0 = torch.tensor([[1, 2]]).float()
    assert w0.shape[1] == d
    model.w.weight.data.copy_(w0)

    X = torch.tensor([[-2, 0, 2], [-1, 1, 3]]).float()
    assert X.shape[0] == d
    assert X.shape[1] == n

    Y = torch.tensor([[0, 1, 2]]).float()
    assert Y.shape[1] == X.shape[1]

    data = X.t()  # PyTorch expects batch dimension first
    target = Y.t()
    assert data.shape[0] == n

    output = model(data)
    # residuals, aka e
    residuals = output - Y.t()

    def compute_loss(residuals_):
        return torch.sum(residuals_ * residuals_) / (2 * n)

    loss = compute_loss(residuals)

    assert loss - 8.83333 < 1e-5, torch.norm(loss) - 8.83333

    # use learning rate 0 to avoid changing parameter vector
    optim_kwargs = dict(lr=0, momentum=0, weight_decay=0, l2_reg=0,
                        bias_correction=False, acc_steps=1,
                        curv_type="Cov", curv_shapes={"Linear": "Kron"},
                        momentum_type="preconditioned", )
    curv_args = dict(damping=1, ema_decay=1)  # todo: damping
    optimizer = SecondOrderOptimizer(model, **optim_kwargs, curv_kwargs=curv_args)

    def backward(last_layer: str) -> Callable:
        """Creates closure that backpropagates either from output layer or from loss layer"""

        def closure() -> Tuple[Optional[torch.Tensor], torch.Tensor]:
            optimizer.zero_grad()
            output = model(data)
            if last_layer == "output":
                output.backward(torch.ones_like(target))
                return None, output
            elif last_layer == 'loss':
                loss = compute_loss(output - target)
                loss.backward()
                return loss, output
            else:
                assert False, 'last layer must be "output" or "loss"'

        return closure

    #    loss = compute_loss(output - Y.t())
    #    loss.backward()

    loss, output = optimizer.step(closure=backward('loss'))
    J = X.t()
    A = model.w.data_input
    B = model.w.grad_output * n
    G = residuals.repeat(1, d) * J
    losses = torch.stack([compute_loss(r) for r in residuals])
    g = G.sum(dim=0) / n
    efisher = G.t() @ G / n
    sigma = efisher - u.outer(g, g)
    loss2 = (residuals * residuals).sum() / (2 * n)
    H = J.t() @ J / n
    noise_variance = torch.trace(H.inverse() @ sigma)

    # H is not quite symmetric, make it so
    H = H + H.t()

    # Slow way
    p_sigma = u.lyapunov_lstsq(H, sigma)

    sigma0 = u.to_numpy(sigma)
    H0 = u.to_numpy(H)

    # Alternative faster way
    p_sigma2 = scipy.linalg.solve_lyapunov(H0, sigma0)
    print(f"Error 1: {np.max(abs(H0 @ p_sigma2 + p_sigma2 @ H0 - sigma0))}")
    u.check_close(p_sigma, p_sigma2)

    # alternative through SVD
    p_sigma3 = lyapunov_svd(torch.tensor(H0), torch.tensor(sigma0))
    u.check_close(p_sigma2, p_sigma3)

    # alternative through evals
    p_sigma4 = u.lyapunov_spectral(torch.tensor(H0), torch.tensor(sigma0))
    u.check_close(p_sigma2, p_sigma4)
//  [python]
def test_lyapunov():
    """Test that scipy lyapunov solver works correctly."""
    d = 2
    n = 3
    torch.set_default_dtype(torch.float32)

    model = Net(d)

    w0 = torch.tensor([[1, 2]]).float()
    assert w0.shape[1] == d
    model.w.weight.data.copy_(w0)

    X = torch.tensor([[-2, 0, 2], [-1, 1, 3]]).float()
    assert X.shape[0] == d
    assert X.shape[1] == n

    Y = torch.tensor([[0, 1, 2]]).float()
    assert Y.shape[1] == X.shape[1]

    data = X.t()  # PyTorch expects batch dimension first
    target = Y.t()
    assert data.shape[0] == n

    output = model(data)
    # residuals, aka e
    residuals = output - Y.t()

    def compute_loss(residuals_):
        return torch.sum(residuals_ * residuals_) / (2 * n)

    loss = compute_loss(residuals)

    assert loss - 8.83333 < 1e-5, torch.linalg.norm(loss) - 8.83333

    # use learning rate 0 to avoid changing parameter vector
    optim_kwargs = dict(lr=0, momentum=0, weight_decay=0, l2_reg=0,
                        bias_correction=False, acc_steps=1,
                        curv_type="Cov", curv_shapes={"Linear": "Kron"},
                        momentum_type="preconditioned", )
    curv_args = dict(damping=1, ema_decay=1)  # todo: damping
    optimizer = SecondOrderOptimizer(model, **optim_kwargs, curv_kwargs=curv_args)

    def backward(last_layer: str) -> Callable:
        """Creates closure that backpropagates either from output layer or from loss layer"""

        def closure() -> Tuple[Optional[torch.Tensor], torch.Tensor]:
            optimizer.zero_grad()
            output = model(data)
            if last_layer == "output":
                output.backward(torch.ones_like(target))
                return None, output
            elif last_layer == 'loss':
                loss = compute_loss(output - target)
                loss.backward()
                return loss, output
            else:
                assert False, 'last layer must be "output" or "loss"'

        return closure

    #    loss = compute_loss(output - Y.t())
    #    loss.backward()

    loss, output = optimizer.step(closure=backward('loss'))
    J = X.t()
    A = model.w.data_input
    B = model.w.grad_output * n
    G = residuals.repeat(1, d) * J
    losses = torch.stack([compute_loss(r) for r in residuals])
    g = G.sum(dim=0) / n
    efisher = G.t() @ G / n
    sigma = efisher - u.outer(g, g)
    loss2 = (residuals * residuals).sum() / (2 * n)
    H = J.t() @ J / n
    noise_variance = torch.trace(H.inverse() @ sigma)

    # H is not quite symmetric, make it so
    H = H + H.t()

    # Slow way
    p_sigma = u.lyapunov_lstsq(H, sigma)

    sigma0 = u.to_numpy(sigma)
    H0 = u.to_numpy(H)

    # Alternative faster way
    p_sigma2 = scipy.linalg.solve_lyapunov(H0, sigma0)
    print(f"Error 1: {np.max(abs(H0 @ p_sigma2 + p_sigma2 @ H0 - sigma0))}")
    u.check_close(p_sigma, p_sigma2)

    # alternative through SVD
    p_sigma3 = lyapunov_svd(torch.tensor(H0), torch.tensor(sigma0))
    u.check_close(p_sigma2, p_sigma3)

    # alternative through evals
    p_sigma4 = u.lyapunov_spectral(torch.tensor(H0), torch.tensor(sigma0))
    u.check_close(p_sigma2, p_sigma4)
//  [python]
def forward(input):
        eps = torch.finfo(input.dtype).eps
        v, w = input[..., :3], input[..., 3:]
        v_norm = torch.norm(v, 2, dim=-1, keepdim=True)
        w_abs = torch.abs(w)
        v_larger_than_eps = (v_norm > eps)
        w_larger_than_eps = (w_abs > eps)
        idx1 = v_larger_than_eps & w_larger_than_eps
        idx2 = v_larger_than_eps & (~w_larger_than_eps)
        idx3 = (~v_larger_than_eps)

        factor = torch.zeros_like(v_norm, requires_grad=False)
        factor[idx1] = 2.0 * torch.atan(v_norm[idx1]/w[idx1]) / v_norm[idx1]
        factor[idx2] = pm(w[idx2]) * torch.pi / v_norm[idx2]
        factor[idx3] = 2.0 * (1.0 / w[idx3] - v_norm[idx3] * v_norm[idx3] / (3 * w[idx3]**3))
        output = factor * v
        return output
//  [python]
def forward(input):
        eps = torch.finfo(input.dtype).eps
        v, w = input[..., :3], input[..., 3:]
        v_norm = torch.linalg.norm(v, 2, dim=-1, keepdim=True)
        w_abs = torch.abs(w)
        v_larger_than_eps = (v_norm > eps)
        w_larger_than_eps = (w_abs > eps)
        idx1 = v_larger_than_eps & w_larger_than_eps
        idx2 = v_larger_than_eps & (~w_larger_than_eps)
        idx3 = (~v_larger_than_eps)

        factor = torch.zeros_like(v_norm, requires_grad=False)
        factor[idx1] = 2.0 * torch.atan(v_norm[idx1]/w[idx1]) / v_norm[idx1]
        factor[idx2] = pm(w[idx2]) * torch.pi / v_norm[idx2]
        factor[idx3] = 2.0 * (1.0 / w[idx3] - v_norm[idx3] * v_norm[idx3] / (3 * w[idx3]**3))
        output = factor * v
        return output
//  [python]
def forward(input):
        theta = torch.norm(input, 2, dim=-1, keepdim=True)
        theta_half, theta2 = 0.5 * theta, theta * theta
        theta4 = theta2 * theta2

        imag_factor = torch.zeros_like(theta, requires_grad=False)
        real_factor = torch.zeros_like(theta, requires_grad=False)
        idx = (theta > torch.finfo(theta.dtype).eps)
        imag_factor[idx] = torch.sin(theta_half[idx]) / theta[idx]
        real_factor[idx] = torch.cos(theta_half[idx])
        imag_factor[~idx] = 0.5 - (1.0/48.0) * theta2[~idx] + (1.0/3840.0) * theta4[~idx]
        real_factor[~idx] = 1.0 - (1.0/8.0) * theta2[~idx] + (1.0/384.0) * theta4[~idx]

        return torch.cat([input * imag_factor, real_factor], -1)
//  [python]
def forward(input):
        theta = torch.linalg.norm(input, 2, dim=-1, keepdim=True)
        theta_half, theta2 = 0.5 * theta, theta * theta
        theta4 = theta2 * theta2

        imag_factor = torch.zeros_like(theta, requires_grad=False)
        real_factor = torch.zeros_like(theta, requires_grad=False)
        idx = (theta > torch.finfo(theta.dtype).eps)
        imag_factor[idx] = torch.sin(theta_half[idx]) / theta[idx]
        real_factor[idx] = torch.cos(theta_half[idx])
        imag_factor[~idx] = 0.5 - (1.0/48.0) * theta2[~idx] + (1.0/3840.0) * theta4[~idx]
        real_factor[~idx] = 1.0 - (1.0/8.0) * theta2[~idx] + (1.0/384.0) * theta4[~idx]

        return torch.cat([input * imag_factor, real_factor], -1)
//  [python]
def test_continuous_time_system(self):
        system = setup_johansson_continuous_time_system1(5)
        x_equilibrium = torch.tensor([0, 0], dtype=system.dtype)
        T = 15.

        def pruner(x):
            return torch.norm(x - x_equilibrium, p=2) < 0.01

        x0_cost_pairs = hybrid_linear_system.generate_cost_to_go_samples(
            system, self.x0_samples, T, lambda x: torch.norm(x, p=2), False,
            x_equilibrium, pruner)
        self.assertIsInstance(x0_cost_pairs, list)
        self.assertGreater(len(x0_cost_pairs), len(self.x0_samples))
        for x0_cost_pair in x0_cost_pairs:
            self.assertFalse(pruner(x0_cost_pair[0]))
//  [python]
def test_continuous_time_system(self):
        system = setup_johansson_continuous_time_system1(5)
        x_equilibrium = torch.tensor([0, 0], dtype=system.dtype)
        T = 15.

        def pruner(x):
            return torch.linalg.norm(x - x_equilibrium, ord=2) < 0.01

        x0_cost_pairs = hybrid_linear_system.generate_cost_to_go_samples(
            system, self.x0_samples, T, lambda x: torch.linalg.norm(x, ord=2), False,
            x_equilibrium, pruner)
        self.assertIsInstance(x0_cost_pairs, list)
        self.assertGreater(len(x0_cost_pairs), len(self.x0_samples))
        for x0_cost_pair in x0_cost_pairs:
            self.assertFalse(pruner(x0_cost_pair[0]))
//  [python]
def sample_directions(n_dims, orthogonalize, cuda):
    if cuda:
        directions = th.cuda.FloatTensor(n_dims, n_dims).normal_(0, 1) + 1e-6
    else:
        directions = th.FloatTensor(n_dims, n_dims).normal_(0, 1) + 1e-6

    if orthogonalize:
        directions, r = th.qr(directions)
        d = th.diag(r, 0)
        ph = d.sign()
        directions *= ph
    directions = th.autograd.Variable(directions, requires_grad=False)
    norm_factors = th.norm(directions, p=2, dim=1, keepdim=True)
    directions = directions / norm_factors
    return directions
//  [python]
def sample_directions(n_dims, orthogonalize, cuda):
    if cuda:
        directions = th.cuda.FloatTensor(n_dims, n_dims).normal_(0, 1) + 1e-6
    else:
        directions = th.FloatTensor(n_dims, n_dims).normal_(0, 1) + 1e-6

    if orthogonalize:
        directions, r = th.qr(directions)
        d = th.diag(r, 0)
        ph = d.sign()
        directions *= ph
    directions = th.autograd.Variable(directions, requires_grad=False)
    norm_factors = th.linalg.norm(directions, ord=2, dim=1, keepdim=True)
    directions = directions / norm_factors
    return directions
//  [python]
def test_discrete_time_system(self):
        system = setup_trecate_discrete_time_system()
        x_equilibrium = torch.tensor([0, 0], dtype=system.dtype)

        def pruner(x):
            return torch.norm(x - x_equilibrium, p=2) < 0.01

        N = 20
        x0_cost_pairs = hybrid_linear_system.generate_cost_to_go_samples(
            system, self.x0_samples, N, lambda x: torch.norm(x, p=2), True,
            x_equilibrium, pruner)
        self.assertIsInstance(x0_cost_pairs, list)
        self.assertGreater(len(x0_cost_pairs), len(self.x0_samples))
        for x0_cost_pair in x0_cost_pairs:
            self.assertFalse(pruner(x0_cost_pair[0]))
//  [python]
def test_discrete_time_system(self):
        system = setup_trecate_discrete_time_system()
        x_equilibrium = torch.tensor([0, 0], dtype=system.dtype)

        def pruner(x):
            return torch.linalg.norm(x - x_equilibrium, ord=2) < 0.01

        N = 20
        x0_cost_pairs = hybrid_linear_system.generate_cost_to_go_samples(
            system, self.x0_samples, N, lambda x: torch.linalg.norm(x, ord=2), True,
            x_equilibrium, pruner)
        self.assertIsInstance(x0_cost_pairs, list)
        self.assertGreater(len(x0_cost_pairs), len(self.x0_samples))
        for x0_cost_pair in x0_cost_pairs:
            self.assertFalse(pruner(x0_cost_pair[0]))
//  [python]
def rxso3_Ws(x):
    rotation, sigma = x[..., :3], x[..., 3]
    theta = torch.norm(rotation, 2, -1)

    A = torch.zeros_like(theta, requires_grad=False)
    B = torch.zeros_like(theta, requires_grad=False)
    C = torch.zeros_like(theta, requires_grad=False)

    sigma_larger = (sigma.abs() > torch.finfo(sigma.dtype).eps)
    theta_larger = (theta > torch.finfo(theta.dtype).eps)
    condition1 = (~sigma_larger) & (~theta_larger)
    condition2 = (~sigma_larger) & theta_larger
    condition3 = sigma_larger & (~theta_larger)
    condition4 = sigma_larger & theta_larger

    scale, sigma2, theta2 = sigma.exp(), sigma * sigma, theta * theta
    theta2_inv = 1.0 / theta2

    # condition1
    C[(~sigma_larger)], A[condition1], B[condition1] = 1.0, 0.5, 1.0 / 6

    # condition2
    theta_c2 = theta[condition2]
    A[condition2] = (1.0 - theta_c2.cos()) * theta2_inv[condition2]
    B[condition2] = (theta_c2 - theta_c2.sin()) / (theta2[condition2] * theta_c2)

    # condition3
    C[sigma_larger] = (scale[sigma_larger] - 1.0) / sigma[sigma_larger]
    sigma_c3, scale_c3, sigma2_c3 = sigma[condition3], scale[condition3], sigma2[condition3]
    A[condition3] = (1.0 + (sigma_c3 - 1.0) * scale_c3) / sigma2_c3
    B[condition3] = (0.5 * sigma2_c3 * scale_c3 + scale_c3 - 1.0 - sigma2_c3 * scale_c3) / (sigma2_c3 * sigma_c3)

    # condition4
    sigma_c4, sigma2_c4, scale_c4 = sigma[condition4], sigma2[condition4], scale[condition4]
    theta_c4, theta2_c4, theta2_inv_c4 = theta[condition4], theta2[condition4], theta2_inv[condition4]
    a_c4, b_c4, c_c4 = scale_c4 * theta_c4.sin(), scale_c4 * theta_c4.cos(), (theta2_c4 + sigma2_c4)
    A[condition4] = (a_c4 * sigma_c4 + (1 - b_c4) * theta_c4) / (theta_c4 * c_c4)
    B[condition4] = (C[condition4] - ((b_c4 - 1) * sigma_c4 + a_c4 * theta_c4) / c_c4) * theta2_inv_c4

    K = vec2skew(rotation)
    A = A.unsqueeze(-1).unsqueeze(-1)
    B = B.unsqueeze(-1).unsqueeze(-1)
    C = C.unsqueeze(-1).unsqueeze(-1)
    I = torch.eye(3, device=x.device, dtype=x.dtype).expand(x.shape[:-1]+(3,3))
    return A * K + B * (K@K) + C * I
//  [python]
def rxso3_Ws(x):
    rotation, sigma = x[..., :3], x[..., 3]
    theta = torch.linalg.norm(rotation, 2, -1)

    A = torch.zeros_like(theta, requires_grad=False)
    B = torch.zeros_like(theta, requires_grad=False)
    C = torch.zeros_like(theta, requires_grad=False)

    sigma_larger = (sigma.abs() > torch.finfo(sigma.dtype).eps)
    theta_larger = (theta > torch.finfo(theta.dtype).eps)
    condition1 = (~sigma_larger) & (~theta_larger)
    condition2 = (~sigma_larger) & theta_larger
    condition3 = sigma_larger & (~theta_larger)
    condition4 = sigma_larger & theta_larger

    scale, sigma2, theta2 = sigma.exp(), sigma * sigma, theta * theta
    theta2_inv = 1.0 / theta2

    # condition1
    C[(~sigma_larger)], A[condition1], B[condition1] = 1.0, 0.5, 1.0 / 6

    # condition2
    theta_c2 = theta[condition2]
    A[condition2] = (1.0 - theta_c2.cos()) * theta2_inv[condition2]
    B[condition2] = (theta_c2 - theta_c2.sin()) / (theta2[condition2] * theta_c2)

    # condition3
    C[sigma_larger] = (scale[sigma_larger] - 1.0) / sigma[sigma_larger]
    sigma_c3, scale_c3, sigma2_c3 = sigma[condition3], scale[condition3], sigma2[condition3]
    A[condition3] = (1.0 + (sigma_c3 - 1.0) * scale_c3) / sigma2_c3
    B[condition3] = (0.5 * sigma2_c3 * scale_c3 + scale_c3 - 1.0 - sigma2_c3 * scale_c3) / (sigma2_c3 * sigma_c3)

    # condition4
    sigma_c4, sigma2_c4, scale_c4 = sigma[condition4], sigma2[condition4], scale[condition4]
    theta_c4, theta2_c4, theta2_inv_c4 = theta[condition4], theta2[condition4], theta2_inv[condition4]
    a_c4, b_c4, c_c4 = scale_c4 * theta_c4.sin(), scale_c4 * theta_c4.cos(), (theta2_c4 + sigma2_c4)
    A[condition4] = (a_c4 * sigma_c4 + (1 - b_c4) * theta_c4) / (theta_c4 * c_c4)
    B[condition4] = (C[condition4] - ((b_c4 - 1) * sigma_c4 + a_c4 * theta_c4) / c_c4) * theta2_inv_c4

    K = vec2skew(rotation)
    A = A.unsqueeze(-1).unsqueeze(-1)
    B = B.unsqueeze(-1).unsqueeze(-1)
    C = C.unsqueeze(-1).unsqueeze(-1)
    I = torch.eye(3, device=x.device, dtype=x.dtype).expand(x.shape[:-1]+(3,3))
    return A * K + B * (K@K) + C * I
//  [python]
def norm_and_var_directions(directions):
    if th.is_tensor(directions):
        directions = th.autograd.Variable(directions, requires_grad=False)
    norm_factors = th.norm(directions, p=2, dim=1, keepdim=True)
    directions = directions / norm_factors
    return directions
//  [python]
def norm_and_var_directions(directions):
    if th.is_tensor(directions):
        directions = th.autograd.Variable(directions, requires_grad=False)
    norm_factors = th.linalg.norm(directions, ord=2, dim=1, keepdim=True)
    directions = directions / norm_factors
    return directions
//  [python]
def _encode_image_prompt(self, prompt, device, num_images_per_prompt, do_classifier_free_guidance):
        r"""
        Encodes the prompt into text encoder hidden states.

        Args:
            prompt (`str` or `List[str]`):
                prompt to be encoded
            device: (`torch.device`):
                torch device
            num_images_per_prompt (`int`):
                number of images that should be generated per prompt
            do_classifier_free_guidance (`bool`):
                whether to use classifier free guidance or not
        """

        def normalize_embeddings(encoder_output):
            embeds = self.image_encoder.vision_model.post_layernorm(encoder_output.last_hidden_state)
            embeds = self.image_encoder.visual_projection(embeds)
            embeds_pooled = embeds[:, 0:1]
            embeds = embeds / torch.norm(embeds_pooled, dim=-1, keepdim=True)
            return embeds

        batch_size = len(prompt) if isinstance(prompt, list) else 1

        # get prompt text embeddings
        image_input = self.image_feature_extractor(images=prompt, return_tensors="pt")
        pixel_values = image_input.pixel_values.to(device).to(self.image_encoder.dtype)
        image_embeddings = self.image_encoder(pixel_values)
        image_embeddings = normalize_embeddings(image_embeddings)

        # duplicate image embeddings for each generation per prompt, using mps friendly method
        bs_embed, seq_len, _ = image_embeddings.shape
        image_embeddings = image_embeddings.repeat(1, num_images_per_prompt, 1)
        image_embeddings = image_embeddings.view(bs_embed * num_images_per_prompt, seq_len, -1)

        # get unconditional embeddings for classifier free guidance
        if do_classifier_free_guidance:
            uncond_images = [np.zeros((512, 512, 3)) + 0.5] * batch_size
            uncond_images = self.image_feature_extractor(images=uncond_images, return_tensors="pt")
            pixel_values = uncond_images.pixel_values.to(device).to(self.image_encoder.dtype)
            negative_prompt_embeds = self.image_encoder(pixel_values)
            negative_prompt_embeds = normalize_embeddings(negative_prompt_embeds)

            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
            seq_len = negative_prompt_embeds.shape[1]
            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)

            # For classifier free guidance, we need to do two forward passes.
            # Here we concatenate the unconditional and conditional embeddings into a single batch
            # to avoid doing two forward passes
            image_embeddings = torch.cat([negative_prompt_embeds, image_embeddings])

        return image_embeddings
//  [python]
def _encode_image_prompt(self, prompt, device, num_images_per_prompt, do_classifier_free_guidance):
        r"""
        Encodes the prompt into text encoder hidden states.

        Args:
            prompt (`str` or `List[str]`):
                prompt to be encoded
            device: (`torch.device`):
                torch device
            num_images_per_prompt (`int`):
                number of images that should be generated per prompt
            do_classifier_free_guidance (`bool`):
                whether to use classifier free guidance or not
        """

        def normalize_embeddings(encoder_output):
            embeds = self.image_encoder.vision_model.post_layernorm(encoder_output.last_hidden_state)
            embeds = self.image_encoder.visual_projection(embeds)
            embeds_pooled = embeds[:, 0:1]
            embeds = embeds / torch.linalg.norm(embeds_pooled, dim=-1, keepdim=True)
            return embeds

        batch_size = len(prompt) if isinstance(prompt, list) else 1

        # get prompt text embeddings
        image_input = self.image_feature_extractor(images=prompt, return_tensors="pt")
        pixel_values = image_input.pixel_values.to(device).to(self.image_encoder.dtype)
        image_embeddings = self.image_encoder(pixel_values)
        image_embeddings = normalize_embeddings(image_embeddings)

        # duplicate image embeddings for each generation per prompt, using mps friendly method
        bs_embed, seq_len, _ = image_embeddings.shape
        image_embeddings = image_embeddings.repeat(1, num_images_per_prompt, 1)
        image_embeddings = image_embeddings.view(bs_embed * num_images_per_prompt, seq_len, -1)

        # get unconditional embeddings for classifier free guidance
        if do_classifier_free_guidance:
            uncond_images = [np.zeros((512, 512, 3)) + 0.5] * batch_size
            uncond_images = self.image_feature_extractor(images=uncond_images, return_tensors="pt")
            pixel_values = uncond_images.pixel_values.to(device).to(self.image_encoder.dtype)
            negative_prompt_embeds = self.image_encoder(pixel_values)
            negative_prompt_embeds = normalize_embeddings(negative_prompt_embeds)

            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
            seq_len = negative_prompt_embeds.shape[1]
            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)

            # For classifier free guidance, we need to do two forward passes.
            # Here we concatenate the unconditional and conditional embeddings into a single batch
            # to avoid doing two forward passes
            image_embeddings = torch.cat([negative_prompt_embeds, image_embeddings])

        return image_embeddings
//  [python]
def becke_tesellation(
            self, x: torch.Tensor, r: torch.Tensor, dm: torch.Tensor,
            i: torch.Tensor
    ) -> torch.Tensor:
        """

        Generates a weight scheme to avoid overlap

        :param x: sample
        :param r: centers
        :param dm: centers distance matrix
        :param i: map sample -> center
        :return:
        """
        device = x.device
        dtype = x.dtype

        nx = x.size()[0]  # number of sample matches. It should match nx
        mx = x.size()[1]  # number of samples per molecule
        nr = r.size()[0]  # number of molecules
        mr = r.size()[1]  # number of atoms per molecule
        if nx != nr:
            raise IOError("batch dims don't not match: {:d} {:d}".format(nx, nr))

        # calculate r-r distance matrix

        # calculate x-r distance matrix
        # dim 1 will be r and dim 2 will be x
        # final dims should be [nr, mx, mr]

        x_exp = x.unsqueeze(2).expand(nr, mx, mr, 3)
        r_exp = r.unsqueeze(1).expand(nr, mx, mr, 3)
        rx_dm = torch.norm(x_exp - r_exp, dim=3)
        del x_exp, r_exp

        # calculate (x-r)-(x-r)T distance vectors

        rx_dm_expanded1 = rx_dm.unsqueeze(3).expand(nr, mx, mr, mr)
        rx_dm_expanded2 = rx_dm.unsqueeze(2).expand(nr, mx, mr, mr)
        xx_dm_expanded = dm.unsqueeze(0)
        mu = (rx_dm_expanded1 - rx_dm_expanded2) / xx_dm_expanded.clamp(min=1e-12)

        del rx_dm_expanded1, rx_dm_expanded2, xx_dm_expanded
        del dm, rx_dm
        # soft-cutoff function

        for _ in range(1, self.softening):
            mu = 0.5 * mu * (3.0 - (mu.pow(2.0)))

        mu = 0.5 * (1 - mu)

        # diagonal fill

        diag_fill = torch.eye(
            mr, mr, device=device, dtype=dtype
        ).unsqueeze(0).unsqueeze(1).expand(nr, mx, mr, mr)

        mu += (diag_fill * 0.5)

        # productory

        mu = mu.prod(dim=3)

        # mu should be now [nr, mx, mr] containing
        # the weights of each nuclei for each sampling point

        i = i.unsqueeze(2)

        v_i = torch.gather(mu, 2, i).squeeze(2)
        v_all = mu.sum(2)
        w = v_i / v_all

        # w should be now [nr, mx] containing the weight of each sample

        return w