//  [python]
def get_flow_ode_output(self, batched_data, offset=None, **kwargs):
        make_masks(batched_data)
        n_graphs = batched_data["x"].size()[0]
        device = batched_data["x"].device
        lig_mask = batched_data["lig_mask"]
        init_pos_clone = batched_data["init_pos"].clone()

        if offset is not None:
            ads_mask = (batched_data["tags"] == 3).unsqueeze(-1)
            cell = batched_data["cell"].transpose(1, 2).float() # B x 3 x 3
            offseted_init_pos = tensor_merge(ads_mask, batched_data["init_pos"] + offset.unsqueeze(1), batched_data["init_pos"]).transpose(1, 2).float() # B x 3 x T
            solv = torch.solve(offseted_init_pos, cell).solution.transpose(1, 2) # B x T x 3
            solv[:, :, 0] %= 1.0
            solv[:, :, 0] %= 1.0
            solv[:, :, 1] %= 1.0
            solv[:, :, 1] %= 1.0
            offseted_init_pos = torch.matmul(solv, cell.transpose(1, 2))
            batched_data["init_pos"] = tensor_merge(ads_mask, offseted_init_pos, init_pos_clone)

        init_center = get_init_center_pos(batched_data)
        batched_data["pos"] -= init_center
        batched_data["init_pos"] -= init_center
        batched_data["pred_pos"] -= init_center

        max_num_moveable = torch.max(batched_data["num_moveable"])
        index = torch.arange(n_graphs, device=device)
        trace_index = torch.arange(3 * max_num_moveable, device=device)

        def score_fn(x_pos, batched_data, t, max_num_moveable):
            pos_clone = batched_data["pos"].clone()
            batched_data["ts"] = torch.ones(n_graphs, device=device, dtype=torch.long).fill_(t)
            batched_data["pos"] = tensor_merge_truncated(lig_mask.unsqueeze(-1), x_pos, pos_clone, max_num_moveable)
            a = self.alphas[t]
            ret, _ = self(batched_data, **kwargs)
            ret = -1.0 * ret / (1.0 - a).sqrt()
            batched_data["pos"] = pos_clone
            return ret[:, :max_num_moveable, :]

        def ode_fn(t, pos_and_p_numpy):
            with torch.set_grad_enabled(True) and torch.autograd.set_detect_anomaly(True):
                beta = self.betas[t]
                pos_and_p = pos_and_p_numpy.clone()
                pos = torch.nn.parameter.Parameter(pos_and_p[:-n_graphs].reshape(n_graphs, max_num_moveable, 3).clone(), requires_grad=True)

                if self.num_epsilon_estimator <= 0:
                    reverse_flow_ode_fn = lambda pos: (-0.5 * beta * score_fn(pos, batched_data, t, max_num_moveable) - \
                        (1.0 - (1.0 - beta).sqrt()) * (pos - batched_data["init_pos"][:, :max_num_moveable])).\
                        masked_fill((~lig_mask[:, :max_num_moveable]).unsqueeze(-1), 0.0)
                    reverse_flow_ode = reverse_flow_ode_fn(pos)
                    jacob = torch.autograd.functional.jacobian(reverse_flow_ode_fn, pos)
                    jacob = jacob.reshape(n_graphs, max_num_moveable * 3, n_graphs, max_num_moveable, 3)
                    jacob = jacob.transpose(1, 2).view(n_graphs, n_graphs, max_num_moveable, 3, max_num_moveable, 3)
                    jacob = jacob[index, index]
                    jacob = jacob.reshape(n_graphs, 3 * max_num_moveable, 3 * max_num_moveable)[:, trace_index, trace_index]
                    trace = torch.sum(jacob, dim=-1)
                else:
                    estimated_traces = []
                    for_estimator = -0.5 * beta * score_fn(pos, batched_data, t, max_num_moveable)
                    reverse_flow_ode = (for_estimator - (1.0 - (1.0 - beta).sqrt()) * (pos - batched_data["init_pos"][:, :max_num_moveable]))
                    for_estimator = for_estimator.masked_fill((~lig_mask[:, :max_num_moveable]).unsqueeze(-1), 0.0)
                    reverse_flow_ode = reverse_flow_ode.masked_fill((~lig_mask[:, :max_num_moveable]).unsqueeze(-1), 0.0)
                    for _ in range(self.num_epsilon_estimator):
                        epsilon = torch.zeros_like(batched_data["pos"][:, :max_num_moveable, :]).normal_()
                        drift_prod_sum = torch.sum(torch.sum(epsilon * for_estimator, dim=[1, 2]))
                        grad = torch.autograd.grad(drift_prod_sum, pos, retain_graph=True)[0]
                        grad = grad.masked_fill((~lig_mask[:, :max_num_moveable]).unsqueeze(-1), 0.0)
                        estimated_trace = torch.sum(grad * epsilon, dim=[1, 2]) - (1.0 - (1.0 - beta).sqrt()) * torch.sum(lig_mask[:, :max_num_moveable] * 3.0, dim=1)
                        estimated_traces.append(estimated_trace.unsqueeze(0))
                    estimated_traces = torch.cat(estimated_traces, dim=0)
                    trace = torch.mean(estimated_traces, dim=0)
            ret = torch.cat([reverse_flow_ode.reshape(-1), trace], dim=0).detach()
            return ret, for_estimator.detach()

        init = torch.cat([batched_data["pred_pos"][:, :max_num_moveable, :].reshape(-1), torch.zeros(n_graphs, dtype=torch.float64, device=device)], dim=0)#.cpu()
        for_estimator_sum = torch.zeros_like(batched_data["pred_pos"][:, :max_num_moveable])
        pos_and_p_numpy = init
        for t in range(self.num_timesteps):
            delta_pos_and_p_numpy, delta_for_estimator = ode_fn(t, pos_and_p_numpy)
            pos_and_p_numpy += delta_pos_and_p_numpy
            for_estimator_sum += delta_for_estimator
        sol = pos_and_p_numpy
        likelihood = sol[-n_graphs:]
        latent_pos = (sol[:-n_graphs].reshape(n_graphs, max_num_moveable, 3) - batched_data["init_pos"][:, :max_num_moveable, :])
        lig_mask = lig_mask.unsqueeze(-1)
        latent_pos = latent_pos * lig_mask[:, :max_num_moveable, :]
        prior_log_p = -0.5 * torch.sum(lig_mask * 3.0, dim=[1, 2]) * np.log(2 * np.pi) - torch.sum((latent_pos ** 2) / 2.0, dim=[1, 2])
        likelihood += prior_log_p

        batched_data["pos"] += init_center
        batched_data["init_pos"] += init_center
        batched_data["pred_pos"] += init_center

        return {
            "persample_likelihood": likelihood,
            "persample_prior_log_p": prior_log_p,
            "latent_pos": latent_pos,
            "sample_size": n_graphs,
        }