include/tenzr/txpr/grad.h¶

// Copyright 2026 The Tenzr Project Authors
//
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef TENZR_TXPR_GRAD_H_
#define TENZR_TXPR_GRAD_H_

#include <concepts>
#include <functional>
#include <memory>
#include <type_traits>
#include <unordered_map>
#include <utility>
#include <vector>

#include "tenzr/frontend/array.h"
#include "tenzr/txpr/txpr.h"

namespace tenzr::txpr {

  std::unordered_map<uint64_t, Array> Grad(const Array& loss, const std::vector<Array>& targets);

  Array Grad(const Array& loss, const Array& target);

  struct VjpResult {
    Array output;
    std::function<std::vector<Array>(const Array&)> pullback;
  };

  struct ValueAndGradResult {
    Array value;
    Array grad;
  };

  VjpResult Vjp(const std::function<Array(const std::vector<Array>&)>& fn, const std::vector<Array>& primals);

  std::function<Array(const Array&)> Grad(const std::function<Array(const Array&)>& fn);

  std::function<ValueAndGradResult(const Array&)> ValueAndGrad(const std::function<Array(const Array&)>& fn);

  template <typename Fn>
    requires std::invocable<Fn, const Array&> && (!std::is_same_v<std::decay_t<Fn>, std::function<Array(const Array&)>>)
  inline std::function<Array(const Array&)> Grad(Fn&& fn) {
    return Grad(std::function<Array(const Array&)>(std::forward<Fn>(fn)));
  }

  template <typename Fn>
    requires std::invocable<Fn, const Array&> && (!std::is_same_v<std::decay_t<Fn>, std::function<Array(const Array&)>>)
  inline std::function<ValueAndGradResult(const Array&)> ValueAndGrad(Fn&& fn) {
    return ValueAndGrad(std::function<Array(const Array&)>(std::forward<Fn>(fn)));
  }

}  // namespace tenzr::txpr

namespace tenzr::autograd {
  using txpr::Grad;
  using txpr::ValueAndGrad;
  using txpr::Vjp;
}  // namespace tenzr::autograd

namespace tenzr {
  using txpr::Grad;
  using txpr::ValueAndGrad;
  using txpr::Vjp;
}  // namespace tenzr

#endif  // TENZR_TXPR_GRAD_H_