src/txpr/custom_vjp.cc¶

// 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.

#include "tenzr/txpr/custom_vjp.h"

#include <stdexcept>
#include <string>
#include <utility>

#include "tenzr/common/error.h"
#include "tenzr/frontend/op_registry.h"
#include "tenzr/txpr/tracer.h"

namespace tenzr::txpr {

  CustomVJPRegistry& CustomVJPRegistry::Instance() {
    static CustomVJPRegistry reg;
    return reg;
  }

  void CustomVJPRegistry::RegisterRule(std::string name, CustomVJPRule rule) {
    if (!OpRegistry::Instance().HasOp(name)) {
      OpDef def;
      def.name = name;
      def.Category = OpCategory::kCustom;
      def.shape_fn = [](const std::vector<Shape>& shapes) { return shapes.empty() ? Shape({}) : shapes[0]; };
      def.dtype_fn = [](const std::vector<DType>& dtypes) { return dtypes.empty() ? DType::kFloat32 : dtypes[0]; };
      // Default lowering: the custom op is an identity marker (the gradient rule is
      // what matters). Ops that need a real forward lowering must register their own.
      def.emit_fn = [](const std::string& result_name, const std::string& result_type, const std::vector<std::string>& input_names,
                       const std::vector<std::string>& input_types) -> std::string {
        TENZR_CHECK(!input_names.empty(), "custom-VJP op requires at least one operand");
        return "    " + result_name + " = stablehlo.convert " + input_names[0] + " : (" + input_types[0] + ") -> " + result_type + "\n";
      };
      OpRegistry::Instance().RegisterOp(std::move(def));
    }
    rules_[std::move(name)] = std::move(rule);
  }

  bool CustomVJPRegistry::HasRule(std::string_view name) const { return rules_.find(std::string(name)) != rules_.end(); }

  const CustomVJPRule& CustomVJPRegistry::GetRule(std::string_view name) const {
    auto it = rules_.find(std::string(name));
    TENZR_CHECK(it != rules_.end(), "CustomVJPRegistry: rule not found for " + std::string(name));
    return it->second;
  }

  CustomVJPFunction::CustomVJPFunction(std::string name, std::function<Array(const Array&)> fn, std::function<std::pair<Array, Array>(const Array&)> fwd,
                                       std::function<Array(const Array&, const Array&)> bwd)
      : name_(std::move(name)), fn_(std::move(fn)) {
    CustomVJPRule rule;
    rule.fwd = std::move(fwd);
    rule.bwd = std::move(bwd);
    CustomVJPRegistry::Instance().RegisterRule(name_, std::move(rule));
  }

  namespace {
    constinit thread_local bool s_grad_tracing_active = false;
  }

  void SetGradTracingActive(bool active) noexcept { s_grad_tracing_active = active; }

  bool IsGradTracingActive() noexcept { return s_grad_tracing_active; }

  Array CustomVJPFunction::operator()(const Array& x) const {
    if (TracerContext::IsActive() || x.IsSymbolic() || IsGradTracingActive()) {
      return Array::CustomOp(name_, {x});
    }
    return fn_(x);
  }

  CustomVJPFunction CustomVjp(std::string name, std::function<Array(const Array&)> fn, std::function<std::pair<Array, Array>(const Array&)> fwd,
                              std::function<Array(const Array&, const Array&)> bwd) {
    return CustomVJPFunction(std::move(name), std::move(fn), std::move(fwd), std::move(bwd));
  }

  void RegisterCustomVjp(const std::string& name, std::function<Array(const Array&, const Array&)> bwd_fn) {
    CustomVJPRule rule;
    rule.bwd = std::move(bwd_fn);
    CustomVJPRegistry::Instance().RegisterRule(name, std::move(rule));
  }

}  // namespace tenzr::txpr