Как вычислить Якобиан RNN Hidden-to-Hidden в Pytorch?

Question

Я разместил этот вопрос на доске обсуждений PyTorch , но, возможно, ТАК поможет мне быстрее. Я хочу вычислить якобиан скрытого состояния для скрытого состояния RNN с одним слоем, оцениваемый в нескольких точках (т. Е. За пакет), и я хочу подтвердить, что делаю это правильно.

Скрытое состояние имеет форму (batch size, num layers=1, hidden state size). Сначала я вычисляю единичные базисные векторы как

    num_basis_vectors = num_layers * hidden_size
    unit_basis_vectors = torch.eye(n=num_basis_vectors).reshape(
        num_basis_vectors, num_layers, hidden_size)

Я затем l oop по единичным базисным векторам (как указано выше):

        jacobian_components = []
        for unit_vector in unit_basis_vectors:
            jacobian_component = torch.autograd.grad(
                outputs=rnn_output_hidden_states,  # shape (batch size, num layers=1, rnn hidden state size) 
                inputs=rnn_input_hidden_states,  # shape (batch size, num layers=1, rnn hidden state size)
                grad_outputs=torch.stack([unit_vector] * batch_size),  # repeat batch_size times
                retain_graph=True)[0]
            jacobian_component = torch.mean(jacobian_component, dim=(1,))
            jacobian_components.append(jacobian_component)
        jacobian = torch.stack(jacobian_components, dim=1)

Это правильно? Как мне go проверить правильность?