Numeric and autodiff helpers

src/ngraph/runtime/utils.cpp

@@ -17,6 +17,7 @@
 #include <cmath>
 
 #include "ngraph/function.hpp"
+#include "ngraph/ops/tuple.hpp"
 #include "ngraph/runtime/call_frame.hpp"
 #include "ngraph/runtime/utils.hpp"
 
@@ -81,21 +82,44 @@ template bool ngraph::runtime::all_close<double>(const std::vector<double>& a,
                                                  double rtol,
                                                  double atol);
 
+ngraph::runtime::FunctionSpec::operator std::shared_ptr<Function>() const
+{
+    return std::make_shared<ngraph::Function>(m_result, m_result_type, m_parameters);
+}
+
+// Returns (dy/(dXs))(C, Xs)
+std::shared_ptr<ngraph::runtime::FunctionSpec>
+    ngraph::runtime::derivative(const std::shared_ptr<ngraph::runtime::FunctionSpec>& f)
+{
+    auto Y = f->get_result();
+    auto Xs = f->get_parameters();
+    auto Y_tv_type = std::dynamic_pointer_cast<const ngraph::TensorViewType>(Y->get_value_type());
+    auto C = std::make_shared<ngraph::op::Parameter>(Y_tv_type->get_element_type(),
+                                                     Y_tv_type->get_shape());
+    std::vector<std::shared_ptr<ngraph::Node>> dYdXs(Xs.size());
+    transform(Xs.begin(), Xs.end(), dYdXs.begin(), [C, Y](const std::shared_ptr<ngraph::Node>& X) {
+        return Y->backwards_derivative(X, C);
+    });
+    auto result = std::make_shared<ngraph::op::Tuple>(dYdXs);
+    std::vector<std::shared_ptr<ngraph::op::Parameter>> args;
+    args.push_back(C);
+    args.insert(args.end(), Xs.begin(), Xs.end());
+    return std::make_shared<ngraph::runtime::FunctionSpec>(result, result->get_value_type(), args);
+}
+
 template <typename ET>
 std::vector<std::shared_ptr<ngraph::runtime::ParameterizedTensorView<ET>>>
     ngraph::runtime::numeric_derivative(
         const std::shared_ptr<runtime::Manager>& manager,
         const std::shared_ptr<runtime::Backend>& backend,
-        const std::shared_ptr<ngraph::Function>& f,
+        const std::shared_ptr<FunctionSpec>& f,
         const std::vector<std::shared_ptr<ngraph::runtime::ParameterizedTensorView<ET>>>& args,
         typename ET::type delta)
 {
     auto y = f->get_result();
 
     Shape y_shape =
-        std::dynamic_pointer_cast<const ngraph::descriptor::TensorView>(y->get_value_type())
-            ->get_tensor_view_type()
-            ->get_shape();
+        std::dynamic_pointer_cast<const ngraph::TensorViewType>(y->get_value_type())->get_shape();
 
     // Check all the shapes
     std::vector<std::shared_ptr<ngraph::runtime::ParameterizedTensorView<ET>>> results;
@@ -107,7 +131,7 @@ std::vector<std::shared_ptr<ngraph::runtime::ParameterizedTensorView<ET>>>
         results.push_back(backend->make_parameterized_tensor_view<ET>(s));
     }
 
-    auto external = manager->compile(f);
+    auto external = manager->compile(*f);
     auto cf = backend->make_call_frame(external);
 
     // ref_y is the function evaluated at the args
@@ -155,7 +179,7 @@ template std::vector<
     ngraph::runtime::numeric_derivative(
         const std::shared_ptr<runtime::Manager>& manager,
         const std::shared_ptr<runtime::Backend>& backend,
-        const std::shared_ptr<ngraph::Function>& f,
+        const std::shared_ptr<FunctionSpec>& f,
         const std::vector<
             std::shared_ptr<ngraph::runtime::ParameterizedTensorView<element::Float32>>>& args,
         element::Float32::type delta);
@@ -165,7 +189,95 @@ template std::vector<
     ngraph::runtime::numeric_derivative(
         const std::shared_ptr<runtime::Manager>& manager,
         const std::shared_ptr<runtime::Backend>& backend,
-        const std::shared_ptr<ngraph::Function>& f,
+        const std::shared_ptr<FunctionSpec>& f,
         const std::vector<
             std::shared_ptr<ngraph::runtime::ParameterizedTensorView<element::Float64>>>& args,
         element::Float64::type delta);
+
+template <typename ET>
+std::vector<std::shared_ptr<ngraph::runtime::ParameterizedTensorView<ET>>>
+    ngraph::runtime::backwards_derivative(
+        const std::shared_ptr<runtime::Manager>& manager,
+        const std::shared_ptr<runtime::Backend>& backend,
+        const std::shared_ptr<FunctionSpec>& f,
+        const std::vector<std::shared_ptr<ngraph::runtime::ParameterizedTensorView<ET>>>& args)
+{
+    auto y = f->get_result();
+    Shape y_shape =
+        std::dynamic_pointer_cast<const ngraph::TensorViewType>(y->get_value_type())->get_shape();
+
+    auto c_param = std::make_shared<op::Parameter>(ET::element_type(), y_shape);
+    auto c_arg = backend->make_parameterized_tensor_view<ET>(y_shape);
+    auto params = f->get_parameters();
+
+    std::vector<std::shared_ptr<ngraph::Node>> deriv_nodes;
+    std::vector<std::shared_ptr<ngraph::runtime::ParameterizedTensorView<ET>>> bprops;
+    std::vector<std::shared_ptr<ngraph::runtime::ParameterizedTensorView<ET>>> results;
+    for (auto param : params)
+    {
+        Shape s = y_shape;
+        auto param_shape =
+            std::dynamic_pointer_cast<const ngraph::TensorViewType>(param->get_value_type())
+                ->get_shape();
+        s.insert(s.end(), param_shape.begin(), param_shape.end());
+        results.push_back(backend->make_parameterized_tensor_view<ET>(s));
+        bprops.push_back(backend->make_parameterized_tensor_view<ET>(param_shape));
+        deriv_nodes.push_back(y->backwards_derivative(param, c_param));
+    }
+
+    std::vector<std::shared_ptr<op::Parameter>> df_params = params;
+    df_params.push_back(c_param);
+    auto df_result = std::make_shared<op::Tuple>(deriv_nodes);
+    auto df = std::make_shared<ngraph::Function>(df_result, df_result->get_value_type(), df_params);
+
+    auto external = manager->compile(df);
+    auto cf = backend->make_call_frame(external);
+
+    // We compute the derivatives chunk by chunk
+    std::vector<typename std::vector<typename ET::type>::iterator> result_pos;
+    for (auto result : results)
+    {
+        result_pos.push_back(result->get_vector().begin());
+    }
+
+    ngraph::runtime::TensorViewPtrs args_tv;
+    args_tv.insert(args_tv.begin(), args.begin(), args.end());
+    args_tv.push_back(c_arg);
+
+    TensorViewPtrs bprops_tv;
+    bprops_tv.insert(bprops_tv.begin(), bprops.begin(), bprops.end());
+
+    auto c_vec = c_arg->get_vector();
+    for (size_t i = 0; i < c_vec.size(); i++)
+    {
+        c_vec[i] = 1;
+        cf->tensor_call(args_tv, bprops_tv);
+        c_vec[i] = 0;
+        for (size_t j = 0; j < results.size(); j++)
+        {
+            auto bprop_vec = bprops[j]->get_vector();
+            result_pos[j] =
+                results[j]->get_vector().insert(result_pos[j], bprop_vec.begin(), bprop_vec.end());
+        }
+    }
+
+    return results;
+}
+
+template std::vector<
+    std::shared_ptr<ngraph::runtime::ParameterizedTensorView<ngraph::element::Float32>>>
+    ngraph::runtime::backwards_derivative<ngraph::element::Float32>(
+        const std::shared_ptr<runtime::Manager>& manager,
+        const std::shared_ptr<runtime::Backend>& backend,
+        const std::shared_ptr<FunctionSpec>& f,
+        const std::vector<
+            std::shared_ptr<ngraph::runtime::ParameterizedTensorView<element::Float32>>>& args);
+
+template std::vector<
+    std::shared_ptr<ngraph::runtime::ParameterizedTensorView<ngraph::element::Float64>>>
+    ngraph::runtime::backwards_derivative<ngraph::element::Float64>(
+        const std::shared_ptr<runtime::Manager>& manager,
+        const std::shared_ptr<runtime::Backend>& backend,
+        const std::shared_ptr<FunctionSpec>& f,
+        const std::vector<
+            std::shared_ptr<ngraph::runtime::ParameterizedTensorView<element::Float64>>>& args);

src/ngraph/runtime/utils.hpp

@@ -90,6 +90,44 @@ namespace ngraph
                                                                 double rtol,
                                                                 double atol);
 
+        class FunctionSpec
+        {
+        public:
+            FunctionSpec(const std::shared_ptr<Node>& result,
+                         const std::shared_ptr<const ValueType>& result_type,
+                         const std::vector<std::shared_ptr<op::Parameter>>& parameters)
+                : m_result(result)
+                , m_result_type(result_type)
+                , m_parameters(parameters)
+            {
+            }
+
+            FunctionSpec(const std::shared_ptr<Node>& result,
+                         const std::vector<std::shared_ptr<op::Parameter>>& parameters)
+                : m_result(result)
+                , m_result_type(result->get_value_type())
+                , m_parameters(parameters)
+            {
+            }
+
+            const std::shared_ptr<const ValueType> get_result_type() const { return m_result_type; }
+            std::shared_ptr<Node> get_result() { return m_result; }
+            const std::vector<std::shared_ptr<op::Parameter>> get_parameters() const
+            {
+                return m_parameters;
+            }
+
+            operator std::shared_ptr<Function>() const;
+
+        protected:
+            std::shared_ptr<Node> m_result;
+            std::shared_ptr<const ValueType> m_result_type;
+            std::vector<std::shared_ptr<op::Parameter>> m_parameters;
+        };
+
+        std::shared_ptr<ngraph::runtime::FunctionSpec>
+            derivative(const std::shared_ptr<ngraph::runtime::FunctionSpec>& f);
+
         /// @brief numeric approximation of the derivative
         /// @param f A function
         /// @param args Values for the arguments (the independent variables)
@@ -100,7 +138,7 @@ namespace ngraph
             numeric_derivative(
                 const std::shared_ptr<runtime::Manager>& manager,
                 const std::shared_ptr<runtime::Backend>& backend,
-                const std::shared_ptr<Function>& f,
+                const std::shared_ptr<FunctionSpec>& f,
                 const std::vector<std::shared_ptr<ngraph::runtime::ParameterizedTensorView<ET>>>&
                     args,
                 typename ET::type delta);
@@ -110,7 +148,7 @@ namespace ngraph
             ngraph::runtime::numeric_derivative(
                 const std::shared_ptr<runtime::Manager>& manager,
                 const std::shared_ptr<runtime::Backend>& backend,
-                const std::shared_ptr<ngraph::Function>& f,
+                const std::shared_ptr<FunctionSpec>& f,
                 const std::vector<
                     std::shared_ptr<ngraph::runtime::ParameterizedTensorView<element::Float32>>>&
                     args,
@@ -121,10 +159,39 @@ namespace ngraph
             ngraph::runtime::numeric_derivative(
                 const std::shared_ptr<runtime::Manager>& manager,
                 const std::shared_ptr<runtime::Backend>& backend,
-                const std::shared_ptr<ngraph::Function>& f,
+                const std::shared_ptr<FunctionSpec>& f,
                 const std::vector<
                     std::shared_ptr<ngraph::runtime::ParameterizedTensorView<element::Float64>>>&
                     args,
                 element::Float64::type delta);
+
+        template <typename ET>
+        std::vector<std::shared_ptr<ngraph::runtime::ParameterizedTensorView<ET>>>
+            backwards_derivative(
+                const std::shared_ptr<runtime::Manager>& manager,
+                const std::shared_ptr<runtime::Backend>& backend,
+                const std::shared_ptr<FunctionSpec>& f,
+                const std::vector<std::shared_ptr<ngraph::runtime::ParameterizedTensorView<ET>>>&
+                    args);
+
+        extern template std::vector<
+            std::shared_ptr<ngraph::runtime::ParameterizedTensorView<ngraph::element::Float32>>>
+            ngraph::runtime::backwards_derivative<ngraph::element::Float32>(
+                const std::shared_ptr<runtime::Manager>& manager,
+                const std::shared_ptr<runtime::Backend>& backend,
+                const std::shared_ptr<FunctionSpec>& f,
+                const std::vector<
+                    std::shared_ptr<ngraph::runtime::ParameterizedTensorView<element::Float32>>>&
+                    args);
+
+        extern template std::vector<
+            std::shared_ptr<ngraph::runtime::ParameterizedTensorView<ngraph::element::Float64>>>
+            ngraph::runtime::backwards_derivative<ngraph::element::Float64>(
+                const std::shared_ptr<runtime::Manager>& manager,
+                const std::shared_ptr<runtime::Backend>& backend,
+                const std::shared_ptr<FunctionSpec>& f,
+                const std::vector<
+                    std::shared_ptr<ngraph::runtime::ParameterizedTensorView<element::Float64>>>&
+                    args);
     }
 }

test/autodiff.cpp

@@ -46,39 +46,20 @@ TEST(backwards, add)
     ASSERT_EQ(DYDX1, C);
 }
 
-// Returns (dy/(dXs))(C, Xs)
-shared_ptr<Function> derivative(const std::shared_ptr<Node>& Y,
-                                const std::vector<std::shared_ptr<op::Parameter>> Xs)
-{
-    auto Y_tv_type = dynamic_pointer_cast<const TensorViewType>(Y->get_value_type());
-    auto C = make_shared<op::Parameter>(Y_tv_type->get_element_type(), Y_tv_type->get_shape());
-    std::vector<std::shared_ptr<Node>> dYdXs(Xs.size());
-    transform(Xs.begin(), Xs.end(), dYdXs.begin(), [C, Y](const std::shared_ptr<Node>& X) {
-        return Y->backwards_derivative(X, C);
-    });
-    auto result = make_shared<op::Tuple>(dYdXs);
-    std::vector<std::shared_ptr<op::Parameter>> args;
-    args.push_back(C);
-    args.insert(args.end(), Xs.begin(), Xs.end());
-    return make_shared<Function>(result, result->get_value_type(), args);
-}
-
 TEST(backwards, multiply)
 {
     auto shape = Shape{2, 3};
     auto make_graph = [shape]() {
         auto X0 = make_shared<op::Parameter>(element::Float32::element_type(), shape);
         auto X1 = make_shared<op::Parameter>(element::Float32::element_type(), shape);
-        return std::make_tuple(X0 * X1, std::vector<std::shared_ptr<op::Parameter>>{X0, X1});
+        return make_shared<ngraph::runtime::FunctionSpec>(
+            X0 * X1, std::vector<std::shared_ptr<op::Parameter>>{X0, X1});
     };
-    auto val_and_vars = make_graph();
-    auto Y = std::get<0>(val_and_vars);
-    auto vars = std::get<1>(val_and_vars);
-    auto f = derivative(Y, vars);
 
     auto manager = runtime::Manager::get("NGVM");
-    auto external = manager->compile(f);
     auto backend = manager->allocate_backend();
+
+    auto external = manager->compile(*derivative(make_graph()));
     auto cf = backend->make_call_frame(external);
 
     auto x0 = backend->make_parameterized_tensor_view<element::Float32>(
@@ -107,4 +88,18 @@ TEST(backwards, multiply)
         ASSERT_EQ(dx0->get_vector(), dx0_correct);
         ASSERT_EQ(dx1->get_vector(), dx1_correct);
     }
+
+    auto f_num = make_graph();
+    auto results_num =
+        runtime::numeric_derivative<element::Float32>(manager, backend, f_num, {x0, x1}, .01f);
+    auto f_sym = make_graph();
+    auto results_sym =
+        runtime::backwards_derivative<element::Float32>(manager, backend, f_sym, {x0, x1});
+    for (size_t i = 0; i < results_num.size(); ++i)
+    {
+        auto result_num = results_num[i];
+        auto result_sym = results_sym[i];
+        bool ac  = all_close(result_num, result_sym);
+        EXPECT_TRUE(ac);
+    }
 }