{% import "common.j2" as common %}
// Auto Generated, DO NOT EDIT!

#include "paddle/fluid/primitive/rule/vjp/generated/generated_vjp.h"
#include "paddle/fluid/ir/dialect/paddle_dialect/ir/pd_api.h"
#include "paddle/fluid/ir/dialect/paddle_dialect/ir/pd_attribute.h"
#include "paddle/fluid/prim/utils/static/static_global_utils.h"
#include "paddle/fluid/primitive/backend/backend.h"
#include "paddle/fluid/primitive/rule/vjp/details.h"
#include "paddle/fluid/primitive/type/lazy_tensor.h"
#include "paddle/fluid/primitive/utils/utils.h"
#include "paddle/ir/core/operation.h"
#include "paddle/phi/core/flags.h"

PHI_DECLARE_string(tensor_operants_mode);

namespace paddle {
namespace primitive {
{% macro sig(fwd_name, name, inputs, attrs, outputs) -%}
std::vector<std::vector<paddle::Tensor>> {{fwd_name}}_vjp({{common.params(inputs, attrs, attrs is exist_mutable_attribute)}}, const std::vector<std::vector<bool>>& stop_gradients)
{%- endmacro -%}

{% macro body(api) %}
std::vector<std::vector<paddle::Tensor>> vjp_res;
for (auto arg: stop_gradients) {
  vjp_res.push_back(std::vector<paddle::Tensor>(arg.size()));
}
  {% if 'composite' in api and api.name in vjp_comp_white_list %}
if (paddle::prim::StaticCompositeContext::Instance().IsBwdPrimEnabled()) {
{% filter indent(2, True) %}{{body_prim(api)}}{% endfilter %}
} else {
{% filter indent(2, True) %}{{body_unprim(api)}}{% endfilter %}
}
  {% else %}
{{body_unprim(api)}}
  {% endif %}
return vjp_res;
{% endmacro %}

{% macro get_mutable_attribute(attrs, api_name) %}
  {% for i in attrs %}
    {%- if i is mutable_attribute -%}
auto* {{i.name}}_define_op = std::static_pointer_cast<primitive::LazyTensor>({{i.name~'_'}}.impl())->getValue().dyn_cast<ir::OpResult>().GetDefiningOp();
      {% if i.typename is scalar %}
if({{i.name}}_define_op->name() != "pd.full") {
  PADDLE_THROW(platform::errors::Unimplemented(
      "We don't support dynamic tensors attribute {{i.name}} for {{api_name}} composite "
      "for now. "));
}
auto {{i.name}} = {{i.name}}_define_op->attribute("value").dyn_cast<paddle::dialect::ScalarAttribute>().data();
      {% elif i.typename is intarray %}
if({{i.name}}_define_op->name() != "pd.full_int_array"){
  PADDLE_THROW(platform::errors::Unimplemented(
      "We don't support dynamic tensors attribute {{i.name}} for {{api_name}} composite "
      "for now. "));
}
auto {{i.name}} = {{i.name}}_define_op->attribute("value").dyn_cast<paddle::dialect::IntArrayAttribute>().data();
      {% endif %}
    {% endif %}
  {% endfor %}
{% endmacro %}

{% macro body_unprim(api) %}
  {%- set input_names=[] -%}
  {%- for i in api.inputs -%} {%- do input_names.append(i.name) -%} {%- endfor -%}
  {%- set attr_names=[] -%}
  {%- for i in api.attrs -%} 
    {%- if i is mutable_attribute -%}
      {%- do input_names.append(i.name~'_') -%} 
    {%- else -%}
      {%- do attr_names.append(i.name) -%} 
    {%- endif -%}
  {%- endfor %}
auto op_res = backend::{{api.name}}<LazyTensor>({{common.args(input_names, attr_names)}});
  {% set outputs = api.outputs|trip_intermediate %} {#- ignore intermediate output -#}
  {% if outputs|length > 1 %}
    {% for i in range(outputs|length) %}
auto out{{i}} = std::get<{{i}}>(op_res);
      {% if outputs[i].typename=='Tensor' %}
vjp_res[{{i}}][0] = !stop_gradients[{{i}}][0] ? out{{i}} : vjp_res[{{i}}][0];
      {% else %}
for (size_t i=0; i< stop_gradients[{{i}}].size(); i++ ) {
    vjp_res[{{i}}][i] =  !stop_gradients[{{i}}][i] ? out{{i}}[i] : vjp_res[{{i}}][i];
}
      {% endif %}
    {% endfor %}
  {% elif outputs|length == 1 %}
    {% if outputs[0].typename=='Tensor' %}
vjp_res[0][0] = !stop_gradients[0][0] ? op_res : vjp_res[0][0];
    {% else %}
for (size_t i=0; i< stop_gradients[0].size(); i++ ) {
  vjp_res[0][i] =  !stop_gradients[0][i] ? op_res[i] : vjp_res[0][i];
}
    {% endif %}
  {% else %} {#- render nothing -#}
  {% endif %}
{% endmacro %}

{% macro body_prim(api) %}
FLAGS_tensor_operants_mode = "static";
  {% for i in range(api.outputs|length) %}
    {% if api.outputs[i].typename=='Tensor' %}
paddle::Tensor* {{api.outputs[i].name}} = !stop_gradients[{{i}}][0] ? &vjp_res[{{i}}][0] : nullptr; 
    {% else %}
std::vector<paddle::Tensor*> {{api.outputs[i].name}}(stop_gradients[{{i}}].size(), nullptr);
for (size_t i=0; i< stop_gradients[{{i}}].size(); i++ ) {
  {{api.outputs[i].name}} =  !stop_gradients[{{i}}][i] ?  &vjp_res[{{i}}][i] : nullptr;
}
    {% endif %}
  {% endfor %}
{{get_mutable_attribute(api.attrs, api.name)}}
details::{{api.composite.func_name}}<LazyTensor>({{api.composite.func_args}});
{% endmacro %}

{%- set api_map = {} -%}
{%- for api in apis -%} {%- do api_map.update({api.name: api}) -%} {%- endfor -%}
{%- for api in apis %}
  {%- if api.backward and api.backward in api_map and api.backward in vjp_white_list -%}
      {%- set backward_api = api_map[api.backward] %}
{{sig(api.name, backward_api.name, backward_api.inputs, backward_api.attrs, backward_api.outputs)}} {
    {% filter indent(2, True) %}
{{body(backward_api)}}
    {% endfilter %}
}

  {% endif %}
{% endfor %}


}  // namespace primitive
}  // namespace paddle
