What is the tensor shape for input "x" in D3PMCorruption?

[qchempku2017]

Dear mattergen developers,
I am trying to implement Repaint-style inpainting in mattergen. To achieve this, I have to implement one-step-forward sampling in order to use the resampling technique mentioned in the Repaint paper. Currently, I have some issues in D3PMCorruption. Here, I implement it as follows:

    # TODO: pay attention to the usage of batch_idx and x.
    def sample_next_timestep(
            self,
            x: torch.Tensor,
            t: torch.Tensor,
            dt: torch.Tensor,
            batch_idx: B = None,
            batch: Optional[BatchedData] = None,
    ) -> torch.Tensor:
        """
        Sample next step class x(t+1) from q(x_{t+1} | x_t)

        Args:
            x: current token indices (non-zero based)
            t: continuous time
            dt: unused, only here for API consistency
            batch_idx: index of atoms in a structure, if not None.
              In mattergen, when batch_idx is not None, every batch comprise a certain amount of atoms
              rather than structure. (Check SimpleBatchedData.to_data_list()) When batch_idx is not None, then
            batch: unused, only here for API consistency
        Returns:
            new x(t+1)
        """
        # discrete time index
        t_disc = to_discrete_time(t, N=self.N, T=self.T)
        # t + 1, same reason as in marginal_probs.
        t_disc = maybe_expand(t_disc, batch_idx).long() + 1 # shape [batch]

        x0 = self._to_zero_based(x.long())  # shape [batch, num_nodes]?

        # must be matrix based diffusion to allow "get".
        if not hasattr(self.d3pm, "get"):
            raise NotImplementedError(
                "D3PMCorruption requires a matrix-based scheduler with .get(t). "
                "Other diffusions must implement q(x_{t+1}|x_t) explicitly."
            )

        # transition matrix q(x_{t+1}|x_t), shape [batch, dim, dim]
        Q = self.d3pm.get(t_disc)
        # Q to [batch, num_nodes, to_classes, from_classes]
        Q = Q[:, None, :, :].expand(-1, x0.shape[1], -1, -1)


        # Q is [batch, num_nodes, to_classes, from_classes]
        # prob is [batch, num_nodes, to_classes]
        # dim=3 gather gets output[i][j][k][l] = input[i][j][k][index[i][j][k][l]],
        # output shape will be the same as the index.
        probs = Q.gather(
            dim=3,
            index=x0.unsqueeze(-1).unsqueeze(-1).expand(-1, -1, Q.size(2), 1)  # [B, N, to_classes, 1]
        ).squeeze(-1)

        sample = torch.distributions.Categorical(probs=probs).sample()

        return self._to_non_zero_based(sample)

However, it seems that the shape of x0 and Q is not correct, and probs could not be retrieved. I am struggling to understand the BatchedData structure, but made nearly no progress. Could you please briefly introduce how the input x and t should be handled correctly in implementation?

Appreciated!

[qchempku2017]

@danielzuegner

[qchempku2017]

If possible, please also help explain what batch_idx represents. I'm having a hard time trying to understand, while there is little clue in the original code about its meaning.