_miopen_impl Class — pytorch Architecture

Source Code

aten/src/ATen/native/miopen/RNN_miopen.cpp lines 899–923

template<typename hidden_type>
std::pair<Tensor, hidden_type> _miopen_impl(
    const Tensor& input, const Tensor& _batch_sizes, const hidden_type& hidden,
    TensorList params, bool has_biases, miopenRNNMode_t mode,
    int64_t num_layers, double dropout_p, bool train, bool bidirectional) {
    auto [hx, cx] = unpack_hidden(hidden);
    int64_t hidden_size = hx.size(2);

    TORCH_CHECK(_batch_sizes.dim() == 1, "batch_sizes tensor should be 1D");
    TORCH_CHECK(
        _batch_sizes.device().is_cpu(),
        "batch_sizes tensor should be on CPU, but got ",
        _batch_sizes.device());
    IntArrayRef batch_sizes { _batch_sizes.data_ptr<int64_t>(), static_cast<size_t>(_batch_sizes.size(0)) };

    Tensor dropout_state = at::empty({0}, input.options());

    auto miopen_output = at::miopen_rnn(
        input, params, has_biases ? 4 : 2,
        hx, cx, static_cast<int>(mode), hidden_size, num_layers, /*batch_first=*/false,
        dropout_p, train, bidirectional, batch_sizes, dropout_state);

    return {std::get<0>(miopen_output),
        pack_hidden<hidden_type>(std::get<1>(miopen_output), std::get<2>(miopen_output))};
}