include/tenzr/txpr/jvp.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_JVP_H_
#define TENZR_TXPR_JVP_H_

#include <concepts>
#include <functional>
#include <utility>

#include "tenzr/array.h"
#include "tenzr/common/export.h"

namespace tenzr::txpr {
  struct JvpResult {
    Array primal;
    Array tangent;
  };

  struct LinearizeResult {
    Array primal;
    std::function<Array(const Array&)> linear;
  };

  [[nodiscard]] TENZR_API JvpResult Jvp(const std::function<Array(const Array&)>& fn, const Array& primal, const Array& tangent);

  template <typename Fn>
    requires std::invocable<Fn, const Array&> && (!std::is_same_v<std::decay_t<Fn>, std::function<Array(const Array&)>>)
  [[nodiscard]] inline JvpResult Jvp(Fn&& fn, const Array& primal, const Array& tangent) {
    return Jvp(std::function<Array(const Array&)>(std::forward<Fn>(fn)), primal, tangent);
  }

  [[nodiscard]] TENZR_API LinearizeResult Linearize(const std::function<Array(const Array&)>& fn, const Array& primal);

  template <typename Fn>
    requires std::invocable<Fn, const Array&> && (!std::is_same_v<std::decay_t<Fn>, std::function<Array(const Array&)>>)
  [[nodiscard]] inline LinearizeResult Linearize(Fn&& fn, const Array& primal) {
    return Linearize(std::function<Array(const Array&)>(std::forward<Fn>(fn)), primal);
  }

}  // namespace tenzr::txpr

namespace tenzr {
  using txpr::Jvp;
  using txpr::Linearize;
}  // namespace tenzr

#endif  // TENZR_TXPR_JVP_H_