vuron/adept: Платформа для обучения и вывода нейронных сетей, с API приближенным к PyTorch.

src/dispatch/log_softmax.hpp

18 строк613 B

kolkir

Fix dispatcher calls and add log_softmax test

18 янв 2025, 23:43

18 янв 2025, 23:43ac46045

#pragma once

#include <adept/dispatch/dispatcher.hpp>
#include <adept/tensor.hpp>

namespace adept {

inline Tensor log_softmax_last_dim_fwd(const Tensor& logits) {
  static auto& func = Dispatcher::instance().find("log_softmax_last_dim_fwd");
  return Dispatcher::instance().call<Tensor, const Tensor&>(func, logits);
}

inline Tensor log_softmax_last_dim_bwd(const Tensor& result, const Tensor& out_grad) {
  static auto& func = Dispatcher::instance().find("log_softmax_last_dim_bwd");
  return Dispatcher::instance().call<Tensor, const Tensor&, const Tensor&>(func, result, out_grad);
}

}  // namespace adept