src/frontend/ufunc.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/frontend/ufunc.h"

#include <mutex>

#include "tenzr/common/error.h"

namespace tenzr {

  UFunc::UFunc(std::string name, UFuncOp op, BinaryFn binary_impl) : name_(std::move(name)), op_(op), is_binary_(true), binary_impl_(std::move(binary_impl)) {}

  UFunc::UFunc(std::string name, UFuncOp op, UnaryFn unary_impl) : name_(std::move(name)), op_(op), is_binary_(false), unary_impl_(std::move(unary_impl)) {}

  Tensor UFunc::operator()(const Tensor& a, const Tensor& b) const {
    TENZR_CHECK(is_binary_, "Cannot invoke unary ufunc '" + name_ + "' with two arguments");
    DType target_dtype = PromoteTypes(a.Dtype(), b.Dtype());
    TENZR_CHECK(Shape::CanBroadcast(a.shape(), b.shape()),
                "Cannot broadcast shapes for ufunc '" + name_ + "': " + a.shape().ToString() + " and " + b.shape().ToString());

    Tensor lhs = (a.Dtype() == target_dtype) ? a : a.Astype(target_dtype);
    Tensor rhs = (b.Dtype() == target_dtype) ? b : b.Astype(target_dtype);

    return binary_impl_(lhs, rhs);
  }

  Tensor UFunc::operator()(const Tensor& a) const {
    TENZR_CHECK(!is_binary_, "Cannot invoke binary ufunc '" + name_ + "' with one argument");
    return unary_impl_(a);
  }

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

  UFuncRegistry::UFuncRegistry() { RegisterBuiltins(); }

  void UFuncRegistry::RegisterBuiltins() {
    RegisterUfunc(std::make_shared<UFunc>("add", UFuncOp::kAdd, [](const Tensor& a, const Tensor& b) { return a + b; }));
    RegisterUfunc(std::make_shared<UFunc>("subtract", UFuncOp::kSubtract, [](const Tensor& a, const Tensor& b) { return a - b; }));
    RegisterUfunc(std::make_shared<UFunc>("multiply", UFuncOp::kMultiply, [](const Tensor& a, const Tensor& b) { return a * b; }));
    RegisterUfunc(std::make_shared<UFunc>("divide", UFuncOp::kDivide, [](const Tensor& a, const Tensor& b) { return a / b; }));

    RegisterUfunc(std::make_shared<UFunc>("negative", UFuncOp::kNegative, [](const Tensor& a) { return a.Neg(); }));
    RegisterUfunc(std::make_shared<UFunc>("absolute", UFuncOp::kAbsolute, [](const Tensor& a) { return a.Abs(); }));
    RegisterUfunc(std::make_shared<UFunc>("exponential", UFuncOp::kExponential, [](const Tensor& a) { return a.Exp(); }));
    RegisterUfunc(std::make_shared<UFunc>("logarithm", UFuncOp::kLogarithm, [](const Tensor& a) { return a.Log(); }));
    RegisterUfunc(std::make_shared<UFunc>("square_root", UFuncOp::kSquareRoot, [](const Tensor& a) { return a.Sqrt(); }));
    RegisterUfunc(std::make_shared<UFunc>("sin", UFuncOp::kSin, [](const Tensor& a) { return a.Sin(); }));
    RegisterUfunc(std::make_shared<UFunc>("cos", UFuncOp::kCos, [](const Tensor& a) { return a.Cos(); }));
    RegisterUfunc(std::make_shared<UFunc>("tanh", UFuncOp::kTanh, [](const Tensor& a) { return a.tanh(); }));
    RegisterUfunc(std::make_shared<UFunc>("sigmoid", UFuncOp::kSigmoid, [](const Tensor& a) { return a.sigmoid(); }));
  }

  void UFuncRegistry::RegisterUfunc(std::shared_ptr<UFunc> ufunc) {
    TENZR_CHECK(ufunc != nullptr, "Cannot register null ufunc");
    ufuncs_[ufunc->Name()] = std::move(ufunc);
  }

  std::shared_ptr<UFunc> UFuncRegistry::Find(const std::string& name) const {
    auto it = ufuncs_.find(name);
    if (it != ufuncs_.end()) {
      return it->second;
    }
    return nullptr;
  }

  bool UFuncRegistry::Has(const std::string& name) const { return ufuncs_.find(name) != ufuncs_.end(); }

  Tensor UFuncRegistry::Add(const Tensor& a, const Tensor& b) { return (*Find("add"))(a, b); }

  Tensor UFuncRegistry::Subtract(const Tensor& a, const Tensor& b) { return (*Find("subtract"))(a, b); }

  Tensor UFuncRegistry::Multiply(const Tensor& a, const Tensor& b) { return (*Find("multiply"))(a, b); }

  Tensor UFuncRegistry::Divide(const Tensor& a, const Tensor& b) { return (*Find("divide"))(a, b); }

  Tensor UFuncRegistry::Negative(const Tensor& a) { return (*Find("negative"))(a); }

  Tensor UFuncRegistry::Absolute(const Tensor& a) { return (*Find("absolute"))(a); }

  Tensor UFuncRegistry::Exponential(const Tensor& a) { return (*Find("exponential"))(a); }

  Tensor UFuncRegistry::Logarithm(const Tensor& a) { return (*Find("logarithm"))(a); }

  Tensor UFuncRegistry::SquareRoot(const Tensor& a) { return (*Find("square_root"))(a); }

  Tensor UFuncRegistry::Sin(const Tensor& a) { return (*Find("sin"))(a); }

  Tensor UFuncRegistry::Cos(const Tensor& a) { return (*Find("cos"))(a); }

  Tensor UFuncRegistry::tanh(const Tensor& a) { return (*Find("tanh"))(a); }

  Tensor UFuncRegistry::sigmoid(const Tensor& a) { return (*Find("sigmoid"))(a); }

}  // namespace tenzr