#include "ccv_nnc.h"

#include "ccv_nnc_easy.h"

#include "ccv_nnc_internal.h"

#include "ccv_internal.h"

#include "_ccv_cnnp_model.h"

// MARK - Add-on Functions

static int _ccv_cnnp_model_clip_grad_norm_reduce_norm2(const ccv_nnc_cmd_t cmd, const ccv_nnc_hint_t hint, const int flags, ccv_nnc_tensor_t* const* const inputs, const int input_size, ccv_nnc_tensor_t* const* const outputs, const int output_size, ccv_nnc_stream_context_t* const stream_context)

{

  const int device_id = CCV_TENSOR_GET_DEVICE_ID(inputs[0]->info.type);

  ccv_nnc_tensor_t* const old_norm2 = outputs[1 + device_id * 2];

  ccv_nnc_tensor_t* const norm2 = outputs[1 + device_id * 2 + 1];

  const int tensor_count = ccv_nnc_tensor_count(inputs[0]->info);

  if (tensor_count == 1)

    ccv_nnc_cmd_exec(CMD_MUL_FORWARD(1), hint, flags, TENSOR_LIST(inputs[0], inputs[0]), TENSOR_LIST(norm2), stream_context);

  else {

    ccv_nnc_cmd_exec(CMD_REDUCE_NORM2_FORWARD(), hint, flags, TENSOR_LIST(inputs[0]), TENSOR_LIST(norm2), stream_context);

    ccv_nnc_cmd_exec(CMD_MUL_FORWARD(1), hint, flags, TENSOR_LIST(norm2, norm2), TENSOR_LIST(norm2), stream_context);

  }

  ccv_nnc_cmd_exec(CMD_ADD_FORWARD(1, 1), hint, flags, TENSOR_LIST(old_norm2, norm2), TENSOR_LIST(old_norm2), stream_context);

  return CCV_NNC_EXEC_SUCCESS;

}

static ccv_nnc_cmd_vtab_t clip_grad_norm_reduce_norm2_vtab = {

  .exec = _ccv_cnnp_model_clip_grad_norm_reduce_norm2

};

static int _ccv_cnnp_model_clip_grad_norm_scatter_norm2(const ccv_nnc_cmd_t cmd, const ccv_nnc_hint_t hint, const int flags, ccv_nnc_tensor_t* const* const inputs, const int input_size, ccv_nnc_tensor_t* const* const outputs, const int output_size, ccv_nnc_stream_context_t* const stream_context)

{

  const int device_id = CCV_TENSOR_GET_DEVICE_ID(inputs[0]->info.type);

  ccv_nnc_tensor_t* const norm2 = inputs[1 + device_id * 2];

  ccv_nnc_cmd_exec(CMD_MUL_FORWARD(1), hint, flags, TENSOR_LIST(inputs[0], norm2), TENSOR_LIST(outputs[0]), stream_context);

  return CCV_NNC_EXEC_SUCCESS;

}

static ccv_nnc_cmd_vtab_t clip_grad_norm_scatter_norm2_vtab = {

  .exec = _ccv_cnnp_model_clip_grad_norm_scatter_norm2

};

void ccv_cnnp_model_parameters_clip_grad_norm(ccv_cnnp_model_t* const model, const ccv_cnnp_model_io_t parameters, int norm_type, float max_norm, ccv_nnc_stream_context_t* const stream_context)

{

  assert(norm_type == 2);

  ccv_cnnp_compiled_data_t* const compiled_data = model->compiled_data;

  assert(compiled_data);

  const int parallel_count = ccv_max(model->parallel_count, 1);

  ccv_nnc_tensor_t* norm2[parallel_count * 2];

  ccv_nnc_tensor_t* max_normt[parallel_count];

  const int stream_type = model->compiled_data->stream_type;

  int i;

  if (stream_type == CCV_STREAM_CONTEXT_GPU)

  {

    for (i = 0; i < parallel_count; i++)

    {

      ccv_nnc_tensor_param_t info = {

        .type = CCV_TENSOR_GPU_MEMORY,

        .format = CCV_TENSOR_FORMAT_NHWC,

        .datatype = CCV_32F,

        .dim = {1},

      };

      CCV_TENSOR_SET_DEVICE_ID(info.type, i);

      norm2[i * 2] = ccv_nnc_tensor_new(ccv_nnc_xpu_alloc(&compiled_data->xpu_alloc, i, stream_context, ccv_nnc_tensor_data_size(info)), info, 0);

      norm2[i * 2 + 1] = ccv_nnc_tensor_new(ccv_nnc_xpu_alloc(&compiled_data->xpu_alloc, i, stream_context, ccv_nnc_tensor_data_size(info)), info, 0);

      max_normt[i] = ccv_nnc_tensor_new(ccv_nnc_xpu_alloc(&compiled_data->xpu_alloc, i, stream_context, ccv_nnc_tensor_data_size(info)), info, 0);

    }

  } else {

    for (i = 0; i < parallel_count; 
i++2
)

    {

      ccv_nnc_tensor_param_t info = {

        .type = CCV_TENSOR_CPU_MEMORY,

        .format = CCV_TENSOR_FORMAT_NHWC,

        .datatype = CCV_32F,

        .dim = {1},

      };

      norm2[i * 2] = ccv_nnc_tensor_new(0, info, 0);

      norm2[i * 2 + 1] = ccv_nnc_tensor_new(0, info, 0);

      max_normt[i] = ccv_nnc_tensor_new(0, info, 0);

    }

  }

  // zero out old norm2.

  if (parallel_count > 1)

  {

    ccv_nnc_stream_context_t* streams[parallel_count];

    ccv_nnc_stream_signal_t* signal;

    if (stream_context)

      signal = ccv_nnc_stream_context_emit_signal_new(stream_context);

    for (i = 0; i < parallel_count; i++)

    {

      const int stream_type = CCV_TENSOR_GET_MEMORY(norm2[i * 2]->info.type) == CCV_TENSOR_GPU_MEMORY ? CCV_STREAM_CONTEXT_GPU : CCV_STREAM_CONTEXT_CPU;

      const int device_id = CCV_TENSOR_GET_DEVICE_ID(norm2[i * 2]->info.type);

      int type = stream_type;

      CCV_STREAM_SET_DEVICE_ID(type, device_id);

      ccv_nnc_stream_context_t* const stream_0 = ccv_cnnp_compiled_data_get_stream(compiled_data, type);

      // Wait signal to finish.

      if (stream_context)

        ccv_nnc_stream_context_wait_signal(stream_0, signal);

      ccv_nnc_cmd_exec(CMD_SET_FORWARD(0), ccv_nnc_no_hint, 0, 0, 0, TENSOR_LIST(norm2[i * 2]), stream_0);

      if (stream_context)

      {

        ccv_nnc_stream_signal_t* const signal = ccv_nnc_stream_context_emit_signal_new(stream_0);

        ccv_nnc_stream_context_wait_signal(stream_context, signal);

      }

      streams[i] = stream_0;

    }

    // If this should be blocking, blocking it.

    if (!stream_context)

      for (i = 0; i < parallel_count; i++)

        if (streams[i])

          ccv_nnc_stream_context_wait(streams[i]);

  } else {

    ccv_nnc_cmd_exec(CMD_SET_FORWARD(0), ccv_nnc_no_hint, 0, 0, 0, TENSOR_LIST(norm2[0]), stream_context);

  }

  // Gather norm2.

  ccv_nnc_cmd_t reduce_cmd = {

    .cmd = CCV_NNC_CUSTOM_FORWARD,

    .isa = &clip_grad_norm_reduce_norm2_vtab,

  };

  ccv_cnnp_model_parameter_gradients_map(model, parameters, reduce_cmd, ccv_nnc_no_hint, 0, 0, 0, norm2, parallel_count * 2, stream_context);

  // Now compute max(max_norm / norm2, 1.0).

  if (parallel_count > 1)

  {

    ccv_nnc_stream_context_t* streams[parallel_count];

    ccv_nnc_stream_signal_t* signal;

    if (stream_context)

      signal = ccv_nnc_stream_context_emit_signal_new(stream_context);

    for (i = 0; i < parallel_count; i++)

    {

      const int stream_type = CCV_TENSOR_GET_MEMORY(norm2[i * 2]->info.type) == CCV_TENSOR_GPU_MEMORY ? CCV_STREAM_CONTEXT_GPU : CCV_STREAM_CONTEXT_CPU;

      const int device_id = CCV_TENSOR_GET_DEVICE_ID(norm2[i * 2]->info.type);

      int type = stream_type;

      CCV_STREAM_SET_DEVICE_ID(type, device_id);

      ccv_nnc_stream_context_t* const stream_0 = ccv_cnnp_compiled_data_get_stream(compiled_data, type);

      // Wait signal to finish.

      if (stream_context)

        ccv_nnc_stream_context_wait_signal(stream_0, signal);

      ccv_nnc_cmd_exec(CMD_EWSQRT_FORWARD(), ccv_nnc_no_hint, 0, TENSOR_LIST(norm2[i * 2]), TENSOR_LIST(norm2[i * 2]), stream_0);

      ccv_nnc_cmd_exec(CMD_SET_FORWARD(max_norm), ccv_nnc_no_hint, 0, 0, 0, TENSOR_LIST(max_normt[i]), stream_0);

      ccv_nnc_cmd_exec(CMD_EWDIV_FORWARD(), ccv_nnc_no_hint, 0, TENSOR_LIST(max_normt[i], norm2[i * 2]), TENSOR_LIST(norm2[i * 2]), stream_0);

      ccv_nnc_cmd_exec(CMD_CLAMP_FORWARD(NAN, 1), ccv_nnc_no_hint, 0, TENSOR_LIST(norm2[i * 2]), TENSOR_LIST(norm2[i * 2]), stream_0);

      if (stream_context)

      {

        ccv_nnc_stream_signal_t* const signal = ccv_nnc_stream_context_emit_signal_new(stream_0);

        ccv_nnc_stream_context_wait_signal(stream_context, signal);

      }

      streams[i] = stream_0;

    }

    // If this should be blocking, blocking it.

    if (!stream_context)

      for (i = 0; i < parallel_count; i++)

        if (streams[i])

          ccv_nnc_stream_context_wait(streams[i]);

  } else {

    ccv_nnc_cmd_exec(CMD_EWSQRT_FORWARD(), ccv_nnc_no_hint, 0, TENSOR_LIST(norm2[0]), TENSOR_LIST(norm2[0]), stream_context);

    ccv_nnc_cmd_exec(CMD_SET_FORWARD(max_norm), ccv_nnc_no_hint, 0, 0, 0, TENSOR_LIST(max_normt[0]), stream_context);

    ccv_nnc_cmd_exec(CMD_EWDIV_FORWARD(), ccv_nnc_no_hint, 0, TENSOR_LIST(max_normt[0], norm2[0]), TENSOR_LIST(norm2[0]), stream_context);

    ccv_nnc_cmd_exec(CMD_CLAMP_FORWARD(NAN, 1), ccv_nnc_no_hint, 0, TENSOR_LIST(norm2[0]), TENSOR_LIST(norm2[0]), stream_context);

  }

  ccv_nnc_cmd_t scatter_cmd = {

    .cmd = CCV_NNC_CUSTOM_FORWARD,

    .isa = &clip_grad_norm_scatter_norm2_vtab,

  };

  ccv_cnnp_model_parameter_gradients_map(model, parameters, scatter_cmd, ccv_nnc_no_hint, 0, norm2, parallel_count * 2, 0, 0, stream_context);

  if (stream_type == CCV_STREAM_CONTEXT_GPU)

    for (i = 0; i < parallel_count; i++)

    {

      ccv_nnc_xpu_free(&compiled_data->xpu_alloc, norm2[i * 2]->data.u8);

      ccv_nnc_xpu_free(&compiled_data->xpu_alloc, norm2[i * 2 + 1]->data.u8);

      ccv_nnc_xpu_free(&compiled_data->xpu_alloc, max_normt[i]->data.u8);

    }

  for (i = 0; i < parallel_count; 
i++2
)

  {

    ccv_nnc_tensor_free(norm2[i * 2]);

    ccv_nnc_tensor_free(norm2[i * 2 + 1]);

    ccv_nnc_tensor_free(max_normt[i]);

  }

}

// MARK - Add-on Functions

static int _ccv_cnnp_model_isnan(const ccv_nnc_cmd_t cmd, const ccv_nnc_hint_t hint, const int flags, ccv_nnc_tensor_t* const* const inputs, const int input_size, ccv_nnc_tensor_t* const* const outputs, const int output_size, ccv_nnc_stream_context_t* const stream_context)

{

  const int device_id = CCV_TENSOR_GET_DEVICE_ID(inputs[0]->info.type);

  ccv_nnc_tensor_t* const old_isnanr = outputs[1 + device_id * 2];

  ccv_nnc_tensor_t* const isnanr = outputs[1 + device_id * 2 + 1];

  ccv_nnc_cmd_t reduce_cmd = CMD_REDUCE_ISNAN_FORWARD();

  reduce_cmd.info.reduce.count = ccv_nnc_tensor_nd(inputs[0]->info.dim);

  int i;

  for (i = 0; i < cmd.info.reduce.count; i++)

    reduce_cmd.info.reduce.axis[i] = i;

  ccv_nnc_cmd_exec(reduce_cmd, hint, flags, TENSOR_LIST(inputs[0]), TENSOR_LIST(isnanr), stream_context);

  ccv_nnc_cmd_exec(CMD_EWSUM_FORWARD(), hint, flags, TENSOR_LIST(old_isnanr, isnanr), TENSOR_LIST(old_isnanr), stream_context);

  return CCV_NNC_EXEC_SUCCESS;

}

static ccv_nnc_cmd_vtab_t reduce_isnan_vtab = {

  .exec = _ccv_cnnp_model_isnan

};

int ccv_cnnp_model_parameter_gradients_isnan(ccv_cnnp_model_t* const model, const ccv_cnnp_model_io_t parameters, ccv_nnc_stream_context_t* const stream_context)

{

  ccv_cnnp_compiled_data_t* const compiled_data = model->compiled_data;

  assert(compiled_data);

  const int parallel_count = ccv_max(model->parallel_count, 1);

  ccv_nnc_tensor_t* isnanr[parallel_count * 2];

  const int stream_type = model->compiled_data->stream_type;

  int i;

  if (stream_type == CCV_STREAM_CONTEXT_GPU)

  {

    for (i = 0; i < parallel_count; i++)

    {

      ccv_nnc_tensor_param_t info = {

        .type = CCV_TENSOR_GPU_MEMORY,

        .format = CCV_TENSOR_FORMAT_NHWC,

        .datatype = CCV_32S,

        .dim = {1},

      };

      CCV_TENSOR_SET_DEVICE_ID(info.type, i);

      isnanr[i * 2] = ccv_nnc_tensor_new(ccv_nnc_xpu_alloc(&compiled_data->xpu_alloc, i, stream_context, ccv_nnc_tensor_data_size(info)), info, 0);

      isnanr[i * 2 + 1] = ccv_nnc_tensor_new(ccv_nnc_xpu_alloc(&compiled_data->xpu_alloc, i, stream_context, ccv_nnc_tensor_data_size(info)), info, 0);

    }

  } else {

    for (i = 0; i < parallel_count; i++)

    {

      ccv_nnc_tensor_param_t info = {

        .type = CCV_TENSOR_CPU_MEMORY,

        .format = CCV_TENSOR_FORMAT_NHWC,

        .datatype = CCV_32S,

        .dim = {1},

      };

      isnanr[i * 2] = ccv_nnc_tensor_new(0, info, 0);

      isnanr[i * 2 + 1] = ccv_nnc_tensor_new(0, info, 0);

    }

  }

  // zero out old isnanr.

  if (parallel_count > 1)

  {

    ccv_nnc_stream_context_t* streams[parallel_count];

    ccv_nnc_stream_signal_t* signal;

    if (stream_context)

      signal = ccv_nnc_stream_context_emit_signal_new(stream_context);

    for (i = 0; i < parallel_count; i++)

    {

      const int stream_type = CCV_TENSOR_GET_MEMORY(isnanr[i * 2]->info.type) == CCV_TENSOR_GPU_MEMORY ? CCV_STREAM_CONTEXT_GPU : CCV_STREAM_CONTEXT_CPU;

      const int device_id = CCV_TENSOR_GET_DEVICE_ID(isnanr[i * 2]->info.type);

      int type = stream_type;

      CCV_STREAM_SET_DEVICE_ID(type, device_id);

      ccv_nnc_stream_context_t* const stream_0 = ccv_cnnp_compiled_data_get_stream(compiled_data, type);

      // Wait signal to finish.

      if (stream_context)

        ccv_nnc_stream_context_wait_signal(stream_0, signal);

      ccv_nnc_cmd_exec(CMD_SET_FORWARD(0), ccv_nnc_no_hint, 0, 0, 0, TENSOR_LIST(isnanr[i * 2]), stream_0);

      if (stream_context)

      {

        ccv_nnc_stream_signal_t* const signal = ccv_nnc_stream_context_emit_signal_new(stream_0);

        ccv_nnc_stream_context_wait_signal(stream_context, signal);

      }

      streams[i] = stream_0;

    }

    // If this should be blocking, blocking it.

    if (!stream_context)

      for (i = 0; i < parallel_count; i++)

        if (streams[i])

          ccv_nnc_stream_context_wait(streams[i]);

  } else

    ccv_nnc_cmd_exec(CMD_SET_FORWARD(0), ccv_nnc_no_hint, 0, 0, 0, TENSOR_LIST(isnanr[0]), stream_context);

  // Gather isnanr.

  ccv_nnc_cmd_t reduce_cmd = {

    .cmd = CCV_NNC_CUSTOM_FORWARD,

    .isa = &reduce_isnan_vtab,

  };

  ccv_cnnp_model_parameter_gradients_map(model, parameters, reduce_cmd, ccv_nnc_no_hint, 0, 0, 0, isnanr, parallel_count * 2, stream_context);

  for (i = 0; i < parallel_count; i++)

    ccv_nnc_tensor_free(isnanr[i * 2 + 1]);

  int retval = 0;

  if (stream_type == CCV_TENSOR_GPU_MEMORY)

  {

    ccv_nnc_tensor_param_t info = {

      .type = CCV_TENSOR_CPU_MEMORY,

      .format = CCV_TENSOR_FORMAT_NHWC,

      .datatype = CCV_32S,

      .dim = {1},

    };

    ccv_nnc_tensor_t* checknan = ccv_nnc_tensor_new(0, info, 0);

    for (i = 0; i < parallel_count; i++)

    {

      ccv_nnc_cmd_exec(CMD_DATA_TRANSFER_FORWARD(), ccv_nnc_no_hint, 0, TENSOR_LIST(isnanr[i * 2]), TENSOR_LIST(checknan), 0);

      if (checknan->data.i32[0] > 0)

      {

        retval = 1;

        break;

      }

    }

    ccv_nnc_tensor_free(checknan);

  } else {

    for (i = 0; i < parallel_count; i++)

      if (isnanr[i * 2]->data.i32[0] > 0)

      {

        retval = 1;

        break;

      }

  }

  for (i = 0; i < parallel_count; i++)

    ccv_nnc_tensor_free(isnanr[i * 2]);

  return retval;

}

// MARK - Core Layers

static void _ccv_cnnp_sum_build(ccv_cnnp_model_t* const self, ccv_nnc_symbolic_graph_t* const graph, const ccv_nnc_tensor_symbol_t* const inputs, const int input_size, ccv_nnc_tensor_symbol_t* const outputs, const int output_size)

{

  PRINT(CCV_CLI_VERBOSE, "[cnnp_sum_build] -\n");

  assert(output_size == 1);

  outputs[0] = ccv_nnc_tensor_symbol_new(graph, ccv_nnc_tensor_symbol_params(graph, inputs[0]), 0);

  ccv_nnc_graph_exec_symbol_new(graph, CMD_EWSUM_FORWARD(), inputs, input_size, outputs, output_size, 0);

}

static ccv_cnnp_model_t* _ccv_cnnp_sum_copy(const ccv_cnnp_model_t* const self, void* const context);

static const ccv_cnnp_model_vtab_t ccv_cnnp_sum_isa = {

  .build = _ccv_cnnp_sum_build,

  .copy = _ccv_cnnp_sum_copy,

};

typedef struct {

  ccv_cnnp_model_t super;

  ccv_nnc_tensor_symbol_t output;

} ccv_cnnp_model_sum_t;

ccv_cnnp_model_t* ccv_cnnp_sum(const char* const name)

{

  ccv_cnnp_model_sum_t* const model_sum = (ccv_cnnp_model_sum_t*)cccalloc(1, sizeof(ccv_cnnp_model_sum_t));

  model_sum->super.isa = &ccv_cnnp_sum_isa;

  model_sum->super.input_size = 0;

  model_sum->super.outputs = &model_sum->output;

  model_sum->super.output_size = 1;

  ccv_cnnp_model_copy_name(&model_sum->super, name);

  return (ccv_cnnp_model_t*)model_sum;

}

static ccv_cnnp_model_t* _ccv_cnnp_sum_copy(const ccv_cnnp_model_t* const self, void* const context)

{

  return ccv_cnnp_sum(self->name);

}

typedef struct {

  ccv_cnnp_model_t super;

  int axis;

  ccv_nnc_tensor_symbol_t output;

} ccv_cnnp_model_concat_t;

static void _ccv_cnnp_concat_build(ccv_cnnp_model_t* const super, ccv_nnc_symbolic_graph_t* const graph, const ccv_nnc_tensor_symbol_t* const inputs, const int input_size, ccv_nnc_tensor_symbol_t* const outputs, const int output_size)

{

  const ccv_cnnp_model_concat_t* const self = (const ccv_cnnp_model_concat_t*)super;

  PRINT(CCV_CLI_VERBOSE, "[cnnp_concat_build] 1. -\n");

  assert(output_size == 1);

  ccv_nnc_tensor_param_t output_params = ccv_nnc_tensor_symbol_params(graph, inputs[0]);

  int i, j;

  if (output_params.dim[0] == 0)

    for (i = 1; i < input_size; i++)

    {

      output_params = ccv_nnc_tensor_symbol_params(graph, inputs[i]);

      if (output_params.dim[0] != 0)

        break;

    }

  const int nd = ccv_nnc_tensor_nd(output_params.dim);

  const int axis = self->axis;

  assert(axis < nd);

  output_params.dim[axis] = 0;

  int input_is_contiguous = 1;

  for (i = 0; i < input_size; 
i++8
)

  {

    const ccv_nnc_tensor_param_t input_params = ccv_nnc_tensor_symbol_params(graph, inputs[i]);

    const int input_nd = ccv_nnc_tensor_nd(input_params.dim);

    if (input_nd == 0)

    {

      PRINT(CCV_CLI_VERBOSE, "[cnnp_concat_build] %d. input[%d]: -\n", i + 2, i);

      input_is_contiguous = 0;

      continue;

    }

    if (CCV_CLI_OUTPUT_LEVEL_IS(CCV_CLI_VERBOSE))

    {

      PRINT(CCV_CLI_VERBOSE, "[cnnp_concat_build] %d. input[%d]: (%d", i + 2, i, input_params.dim[0]);

      int i;

      for (i = 1; i < CCV_NNC_MAX_DIM_ALLOC && input_params.dim[i] > 0; i++)

        PRINT(CCV_CLI_VERBOSE, ", %d", input_params.dim[i]);

      PRINT(CCV_CLI_VERBOSE, ")\n");

    }

    assert(input_nd == nd);

    
for (j = 0; 8
j < nd; 
j++8
)

      if (j != axis)

        { assert(input_params.dim[j] == output_params.dim[j]); }

    output_params.dim[axis] += input_params.dim[axis];

  }

  outputs[0] = ccv_nnc_tensor_symbol_new(graph, output_params, 0);

  int ofs[CCV_NNC_MAX_DIM_ALLOC] = {};

  int stride[CCV_NNC_MAX_DIM_ALLOC] = {};

  ccv_nnc_tensor_get_stride(output_params.dim, stride);

  if (input_is_contiguous)

  {

    ccv_nnc_tensor_symbol_t aliases[input_size];

    for (i = 0; i < input_size; 
i++8
)

    {

      const ccv_nnc_tensor_param_t input_params = ccv_nnc_tensor_symbol_params(graph, inputs[i]);

      aliases[i] = ccv_nnc_tensor_symbol_alias_new(graph, outputs[0], ofs, stride, input_params, 0);

      ofs[axis] += input_params.dim[axis];

    }

    // Format transform is more flexible.

    ccv_nnc_graph_exec_symbol_new(graph, CMD_FORMAT_TRANSFORM_FORWARD(), inputs, input_size, aliases, input_size, "concat");

  } else {

    ccv_nnc_tensor_symbol_t aliases[input_size];

    for (i = 0; i < input_size; i++)

    {

      const ccv_nnc_tensor_param_t input_params = ccv_nnc_tensor_symbol_params(graph, inputs[i]);

      if (input_params.dim[0] == 0)

      {

        // Create a new alias anyway, but not going to use it, in this way, the alias count will match during absorb.

        aliases[i] = ccv_nnc_tensor_symbol_alias_new(graph, outputs[0], ofs, stride, input_params, 0);

        continue;

      }

      aliases[i] = ccv_nnc_tensor_symbol_alias_new(graph, outputs[0], ofs, stride, input_params, 0);

      ofs[axis] += input_params.dim[axis];

    }

    // Format transform is more flexible.

    ccv_nnc_graph_exec_symbol_new(graph, CMD_FORMAT_TRANSFORM_FORWARD(), inputs, input_size, aliases, input_size, "concat");

  }

}

static ccv_cnnp_model_t* _ccv_cnnp_concat_copy(const ccv_cnnp_model_t* const self, void* const context);

static const ccv_cnnp_model_vtab_t ccv_cnnp_concat_isa = {

  .build = _ccv_cnnp_concat_build,

  .copy = _ccv_cnnp_concat_copy,

};

ccv_cnnp_model_t* ccv_cnnp_concat(const int axis, const char* const name)

{

  ccv_cnnp_model_concat_t* const model_concat = (ccv_cnnp_model_concat_t*)cccalloc(1, sizeof(ccv_cnnp_model_concat_t));

  model_concat->super.isa = &ccv_cnnp_concat_isa;

  model_concat->super.input_size = 0;

  model_concat->super.outputs = &model_concat->output;

  model_concat->super.output_size = 1;

  model_concat->axis = axis;

  ccv_cnnp_model_copy_name(&model_concat->super, name);

  return (ccv_cnnp_model_t*)model_concat;

}

static ccv_cnnp_model_t* _ccv_cnnp_concat_copy(const ccv_cnnp_model_t* const super, void* const context)

{

  const ccv_cnnp_model_concat_t* const self = (const ccv_cnnp_model_concat_t*)super;

  return ccv_cnnp_concat(self->axis, self->super.name);

}

typedef struct {

  ccv_cnnp_model_t super;

  int axis;

  ccv_nnc_tensor_symbol_t outputs[1];

} ccv_cnnp_model_chunk_t;

static void _ccv_cnnp_chunk_build(ccv_cnnp_model_t* const super, ccv_nnc_symbolic_graph_t* const graph, const ccv_nnc_tensor_symbol_t* const inputs, const int input_size, ccv_nnc_tensor_symbol_t* const outputs, const int output_size)

{

  const ccv_cnnp_model_concat_t* const self = (const ccv_cnnp_model_concat_t*)super;

  PRINT(CCV_CLI_VERBOSE, "[cnnp_chunk_build] 1. axis: %d\n", self->axis);

  assert(input_size == 1);

  const ccv_nnc_tensor_param_t input_params = ccv_nnc_tensor_symbol_params(graph, inputs[0]);

  if (CCV_CLI_OUTPUT_LEVEL_IS(CCV_CLI_VERBOSE))

  {

    PRINT(CCV_CLI_VERBOSE, "[cnnp_chunk_build] 2. input: (%d", input_params.dim[0]);

    int i;

    for (i = 1; i < CCV_NNC_MAX_DIM_ALLOC && input_params.dim[i] > 0; i++)

      PRINT(CCV_CLI_VERBOSE, ", %d", input_params.dim[i]);

    PRINT(CCV_CLI_VERBOSE, ")\n");

  }

  ccv_nnc_tensor_param_t output_params = input_params;

  int i;

  const int nd = ccv_nnc_tensor_nd(output_params.dim);

  const int axis = self->axis;

  assert(axis < nd);

  const int n = self->super.output_size;

  assert(n == output_size);

  assert(output_params.dim[axis] % n == 0);

  output_params.dim[axis] = output_params.dim[axis] / n;

  int ofs[CCV_NNC_MAX_DIM_ALLOC] = {};

  int stride[CCV_NNC_MAX_DIM_ALLOC] = {};

  ccv_nnc_tensor_get_stride(input_params.dim, stride);

  ccv_nnc_tensor_symbol_t to = ccv_nnc_tensor_symbol_alias_to(graph, inputs[0]);

  if (to.d == CCV_NNC_NO_TENSOR_SYMBOL) // If we are not reshape an alias, it is straightforward.

  {

    for (i = 0; i < output_size; 
i++4
)

    {

      outputs[i] = ccv_nnc_tensor_symbol_alias_new(graph, inputs[0], ofs, stride, output_params, 0);

      ofs[axis] += output_params.dim[axis];

    }

  } else {

    // Otherwise, we need to check if it is permute. For permute, we cannot do alias directly.

    // We need to first materialize the permute and then run reshape on top of it, otherwise it will be wrong.

    int old_stride[CCV_NNC_MAX_DIM_ALLOC];

    ccv_nnc_tensor_symbol_alias_params(graph, inputs[0], 0, old_stride);

    // We identify permute by checking if the stride is not in descending order.

    // This also covered "permute" through reshape, rather than using ccv_cnnp_permute directly.

    int i, no_permute = 1;

    for (i = 1; no_permute && i < nd; i++)

      if (old_stride[i - 1] < old_stride[i])

        no_permute = 0;

    if (no_permute)

    { // Just straightforward reshape if there is no no permute.

      for (i = 0; i < output_size; i++)

      {

        outputs[i] = ccv_nnc_tensor_symbol_alias_new(graph, inputs[0], ofs, old_stride, output_params, 0);

        ofs[axis] += output_params.dim[axis];

      }

    } else {

      // Otherwise, we first do format transform to plain tensor and then do reshape.

      ccv_nnc_tensor_symbol_t permuted = ccv_nnc_tensor_symbol_new(graph, input_params, 0);

      ccv_nnc_graph_exec_symbol_new(graph, CMD_FORMAT_TRANSFORM_FORWARD(), TENSOR_SYMBOL_LIST(inputs[0]), TENSOR_SYMBOL_LIST(permuted), "reshape");

      for (i = 0; i < output_size; i++)

      {

        outputs[i] = ccv_nnc_tensor_symbol_alias_new(graph, permuted, ofs, stride, output_params, 0);

        ofs[axis] += output_params.dim[axis];

      }

    }

  }

}

static ccv_cnnp_model_t* _ccv_cnnp_chunk_copy(const ccv_cnnp_model_t* const self, void* const context);

static const ccv_cnnp_model_vtab_t ccv_cnnp_chunk_isa = {

  .build = _ccv_cnnp_chunk_build,

  .copy = _ccv_cnnp_chunk_copy,

};

ccv_cnnp_model_t* ccv_cnnp_chunk(const int n, const int axis, const char* const name)

{

  assert(n >= 1);

  ccv_cnnp_model_chunk_t* const model_chunk = (ccv_cnnp_model_chunk_t*)cccalloc(1, sizeof(ccv_cnnp_model_chunk_t) + sizeof(ccv_nnc_tensor_symbol_t) * (n - 1));

  model_chunk->super.isa = &ccv_cnnp_chunk_isa;

  model_chunk->super.input_size = 1;

  model_chunk->super.outputs = model_chunk->outputs;

  model_chunk->super.output_size = n;

  model_chunk->axis = axis;

  ccv_cnnp_model_copy_name(&model_chunk->super, name);

  return (ccv_cnnp_model_t*)model_chunk;

}

static ccv_cnnp_model_t* _ccv_cnnp_chunk_copy(const ccv_cnnp_model_t* const super, void* const context)

{

  const ccv_cnnp_model_chunk_t* const self = (const ccv_cnnp_model_chunk_t*)super;

  return ccv_cnnp_chunk(self->super.output_size, self->axis, self->super.name);

}

typedef struct {

  ccv_cnnp_model_t super;

  ccv_nnc_tensor_symbol_t output;

  int format;

  int dim[CCV_NNC_MAX_DIM_ALLOC];

  int ofs[CCV_NNC_MAX_DIM_ALLOC];

  int stride[CCV_NNC_MAX_DIM_ALLOC];

} ccv_cnnp_model_reshape_t;

static void _ccv_cnnp_reshape_build(ccv_cnnp_model_t* const super, ccv_nnc_symbolic_graph_t* const graph, const ccv_nnc_tensor_symbol_t* const inputs, const int input_size, ccv_nnc_tensor_symbol_t* const outputs, const int output_size)

{

  assert(input_size == 1);

  assert(output_size == 1);

  ccv_cnnp_model_reshape_t* const self = (ccv_cnnp_model_reshape_t*)super;

  if (CCV_CLI_OUTPUT_LEVEL_IS(CCV_CLI_VERBOSE))

  {

    PRINT(CCV_CLI_VERBOSE, "[cnnp_reshape_build] 1. dim: (%d", self->dim[0]);

    int i;

    for (i = 1; i < CCV_NNC_MAX_DIM_ALLOC && self->dim[i] > 0; i++)

      PRINT(CCV_CLI_VERBOSE, ", %d", self->dim[i]);

    const int count = i;

    PRINT(CCV_CLI_VERBOSE, "), ofs: (%d", self->ofs[0]);

    for (i = 1; i < count; i++)

      PRINT(CCV_CLI_VERBOSE, ", %d", self->ofs[i]);

    PRINT(CCV_CLI_VERBOSE, "), stride: (%d", self->stride[0]);

    for (i = 1; i < count; i++)

      PRINT(CCV_CLI_VERBOSE, ", %d", self->stride[i]);

    PRINT(CCV_CLI_VERBOSE, ")\n");

  }

  ccv_nnc_tensor_param_t params = ccv_nnc_tensor_symbol_params(graph, inputs[0]);

  int dim[CCV_NNC_MAX_DIM_ALLOC];

  memcpy(dim, self->dim, sizeof(dim));

  int i, auto_idx = -1;

  size_t known = 1;

  const size_t tensor_count = ccv_nnc_tensor_count(params);

  for (i = 0; i < CCV_NNC_MAX_DIM_ALLOC && dim[i]; 
i++1.19k
)

    if (dim[i] == -1)

      auto_idx = i;

    else

      known *= dim[i];

  if (auto_idx >= 0)

  {

    assert(known > 0 && tensor_count % known == 0);

    dim[auto_idx] = tensor_count / known;

  }

  if (CCV_CLI_OUTPUT_LEVEL_IS(CCV_CLI_VERBOSE))

  {

    PRINT(CCV_CLI_VERBOSE, "[cnnp_reshape_build] 2. input: (%d", params.dim[0]);

    int i;

    for (i = 1; i < CCV_NNC_MAX_DIM_ALLOC && params.dim[i] > 0; i++)

      PRINT(CCV_CLI_VERBOSE, ", %d", params.dim[i]);

    PRINT(CCV_CLI_VERBOSE, ")\n");

  }

  assert(ccv_nnc_dimension_count(dim) <= ccv_nnc_tensor_count(params));

  ccv_nnc_tensor_symbol_t to = ccv_nnc_tensor_symbol_alias_to(graph, inputs[0]);

  int stride_from_dim[CCV_NNC_MAX_DIM_ALLOC];

  if (to.d == CCV_NNC_NO_TENSOR_SYMBOL) // If we are not reshape an alias, it is straightforward.

  {

    memcpy(params.dim, dim, sizeof(params.dim));

    int* stride;

    if (self->stride[0] == 0)

    {

      ccv_nnc_tensor_get_stride(dim, stride_from_dim);

      stride = stride_from_dim;

    } else

      stride = self->stride;

    if (self->format > 0)

      params.format = self->format;

    outputs[0] = ccv_nnc_tensor_symbol_alias_new(graph, inputs[0], self->ofs, stride, params, 0);

  } else {

    // Otherwise, we need to check if it is permute. For permute, we cannot do alias directly.

    // We need to first materialize the permute and then run reshape on top of it, otherwise it will be wrong.

    int old_stride[CCV_NNC_MAX_DIM_ALLOC];

    ccv_nnc_tensor_symbol_alias_params(graph, inputs[0], 0, old_stride);

    // We identify permute by checking if the stride is not in descending order.

    // This also covered "permute" through reshape, rather than using ccv_cnnp_permute directly.

    const int nd = ccv_nnc_tensor_nd(params.dim);

    const int new_nd = ccv_nnc_tensor_nd(dim);

    int i, no_permute = 1;

    // If the new dim has different nd, or we actually have a stride, we need to check if it is no permute or not.

    if (new_nd != nd || 
(0
self->stride[0] != 00
 && 
memcmp(self->stride, old_stride, sizeof(self->stride))0
))

      
for (i = 1; 1
no_permute && 
i < nd1
; 
i++1
)

        if (old_stride[i - 1] < old_stride[i])

          no_permute = 0;

    if (no_permute)

    { // Just straightforward reshape if there is no no permute.

      memcpy(params.dim, dim, sizeof(params.dim));

      int* stride;

      if (self->stride[0] == 0)

      {

        if (new_nd != nd) // Cannot use old stride.

        {

          ccv_nnc_tensor_get_stride(dim, stride_from_dim);

          stride = stride_from_dim;

        } else

          stride = old_stride;

      } else

        stride = self->stride;

      if (self->format > 0)

        params.format = self->format;

      outputs[0] = ccv_nnc_tensor_symbol_alias_new(graph, inputs[0], self->ofs, stride, params, 0);

    } else {

      // Otherwise, we first do format transform to plain tensor and then do reshape.

      ccv_nnc_tensor_symbol_t permuted = ccv_nnc_tensor_symbol_new(graph, params, 0);

      ccv_nnc_graph_exec_symbol_new(graph, CMD_FORMAT_TRANSFORM_FORWARD(), TENSOR_SYMBOL_LIST(inputs[0]), TENSOR_SYMBOL_LIST(permuted), "reshape");

      memcpy(params.dim, dim, sizeof(params.dim));

      int* stride;

      if (self->stride[0] == 0)

      {

        ccv_nnc_tensor_get_stride(dim, stride_from_dim);

        stride = stride_from_dim;

      } else

        stride = self->stride;

      if (self->format > 0)

        params.format = self->format;

      // And then we create alias against the permuted one.

      outputs[0] = ccv_nnc_tensor_symbol_alias_new(graph, permuted, self->ofs, stride, params, 0);

    }

  }

}

static ccv_cnnp_model_t* _ccv_cnnp_reshape_copy(const ccv_cnnp_model_t* const super, void* const context);

static const ccv_cnnp_model_vtab_t ccv_cnnp_reshape_isa = {

  .build = _ccv_cnnp_reshape_build,

  .copy = _ccv_cnnp_reshape_copy,

};

ccv_cnnp_model_t* ccv_cnnp_reshape(const int format, const int dim[CCV_NNC_MAX_DIM_ALLOC], const int ofs[CCV_NNC_MAX_DIM_ALLOC], const int stride[CCV_NNC_MAX_DIM_ALLOC], const char* const name)

{

  ccv_cnnp_model_reshape_t* const model_reshape = (ccv_cnnp_model_reshape_t*)cccalloc(1, sizeof(ccv_cnnp_model_reshape_t));

  model_reshape->super.isa = &ccv_cnnp_reshape_isa;

  model_reshape->super.input_size = 1;

  model_reshape->super.outputs = &model_reshape->output;

  model_reshape->super.output_size = 1;

  ccv_cnnp_model_copy_name(&model_reshape->super, name);

  model_reshape->format = format;

  memcpy(model_reshape->dim, dim, sizeof(model_reshape->dim));

  memcpy(model_reshape->ofs, ofs, sizeof(model_reshape->ofs));

  if (stride[0] != 0)

    memcpy(model_reshape->stride, stride, sizeof(model_reshape->stride));

  return (ccv_cnnp_model_t*)model_reshape;

}

static ccv_cnnp_model_t* _ccv_cnnp_reshape_copy(const ccv_cnnp_model_t* const super, void* const context)

{

  const ccv_cnnp_model_reshape_t* const self = (const ccv_cnnp_model_reshape_t*)super;

  return ccv_cnnp_reshape(self->format, self->dim, self->ofs, self->stride, self->super.name);

}

typedef struct {

  ccv_cnnp_model_t super;

  ccv_nnc_tensor_symbol_t output;

  int type;

  int begin[CCV_NNC_MAX_DIM_ALLOC];

  int end[CCV_NNC_MAX_DIM_ALLOC];

} ccv_cnnp_model_pad_t;

static void _ccv_cnnp_pad_build(ccv_cnnp_model_t* const super, ccv_nnc_symbolic_graph_t* const graph, const ccv_nnc_tensor_symbol_t* const inputs, const int input_size, ccv_nnc_tensor_symbol_t* const outputs, const int output_size)

{

  assert(input_size == 1);

  assert(output_size == 1);

  ccv_cnnp_model_pad_t* const self = (ccv_cnnp_model_pad_t*)super;

  PRINT(CCV_CLI_VERBOSE, "[cnnp_pad_build] -\n");

  const ccv_nnc_tensor_param_t input_params = ccv_nnc_tensor_symbol_params(graph, inputs[0]);

  const int nd = ccv_nnc_tensor_nd(input_params.dim);

  ccv_nnc_tensor_param_t params = input_params;

  int i;

  for (i = 0 ; i < nd; 
i++4
)

    params.dim[i] += self->begin[i] + self->end[i];

  const ccv_nnc_tensor_symbol_t padded = ccv_nnc_tensor_symbol_new(graph, params, 0);

  ccv_nnc_cmd_t pad = CMD_PAD_FORWARD(self->type, (), ());

  memcpy(pad.info.size.dim, self->begin, sizeof(pad.info.size.dim));

  memcpy(pad.info.pad.end, self->end, sizeof(pad.info.pad.end));

  ccv_nnc_graph_exec_symbol_new(graph, pad, TENSOR_SYMBOL_LIST(inputs[0]), TENSOR_SYMBOL_LIST(padded), "pad");

  outputs[0] = padded;

}

static ccv_cnnp_model_t* _ccv_cnnp_pad_copy(const ccv_cnnp_model_t* const super, void* const context);

static const ccv_cnnp_model_vtab_t ccv_cnnp_pad_isa = {

  .build = _ccv_cnnp_pad_build,

  .copy = _ccv_cnnp_pad_copy,

};

ccv_cnnp_model_t* ccv_cnnp_pad(const int type, const int begin[CCV_NNC_MAX_DIM_ALLOC], const int end[CCV_NNC_MAX_DIM_ALLOC], const char* const name)

{

  ccv_cnnp_model_pad_t* const model_pad = (ccv_cnnp_model_pad_t*)cccalloc(1, sizeof(ccv_cnnp_model_pad_t));

  model_pad->super.isa = &ccv_cnnp_pad_isa;

  model_pad->super.input_size = 1;

  model_pad->super.outputs = &model_pad->output;

  model_pad->super.output_size = 1;

  ccv_cnnp_model_copy_name(&model_pad->super, name);

  model_pad->type = type;

  memcpy(model_pad->begin, begin, sizeof(model_pad->begin));

  memcpy(model_pad->end, end, sizeof(model_pad->end));

  return (ccv_cnnp_model_t*)model_pad;

}

static ccv_cnnp_model_t* _ccv_cnnp_pad_copy(const ccv_cnnp_model_t* const super, void* const context)

{

  const ccv_cnnp_model_pad_t* const self = (const ccv_cnnp_model_pad_t*)super;

  return ccv_cnnp_pad(self->type, self->begin, self->end, self->super.name);

}

typedef struct {

  ccv_cnnp_model_t super;

  ccv_nnc_tensor_symbol_t output;

} ccv_cnnp_model_identity_t;

static void _ccv_cnnp_identity_build(ccv_cnnp_model_t* const super, ccv_nnc_symbolic_graph_t* const graph, const ccv_nnc_tensor_symbol_t* const inputs, const int input_size, ccv_nnc_tensor_symbol_t* const outputs, const int output_size)

{

  assert(input_size == 1);

  assert(output_size == 1);

  PRINT(CCV_CLI_VERBOSE, "[cnnp_identity_build] -\n");

  outputs[0] = inputs[0];

}

static ccv_cnnp_model_t* _ccv_cnnp_identity_copy(const ccv_cnnp_model_t* const super, void* const context);

static const ccv_cnnp_model_vtab_t ccv_cnnp_identity_isa = {

  .build = _ccv_cnnp_identity_build,

  .copy = _ccv_cnnp_identity_copy,

};

ccv_cnnp_model_t* ccv_cnnp_identity(const char* const name)

{

  ccv_cnnp_model_identity_t* const model_identity = (ccv_cnnp_model_identity_t*)cccalloc(1, sizeof(ccv_cnnp_model_identity_t));

  model_identity->super.isa = &ccv_cnnp_identity_isa;

  model_identity->super.input_size = 1;

  model_identity->super.outputs = &model_identity->output;

  model_identity->super.output_size = 1;

  ccv_cnnp_model_copy_name(&model_identity->super, name);

  return (ccv_cnnp_model_t*)model_identity;

}

static ccv_cnnp_model_t* _ccv_cnnp_identity_copy(const ccv_cnnp_model_t* const super, void* const context)

{

  const ccv_cnnp_model_identity_t* const self = (const ccv_cnnp_model_identity_t*)super;

  return ccv_cnnp_identity(self->super.name);

}

typedef struct {

  ccv_cnnp_model_t super;

  ccv_nnc_tensor_symbol_t output;

  int index[CCV_NNC_MAX_DIM_ALLOC];

} ccv_cnnp_model_permute_t;

static void _ccv_cnnp_permute_build(ccv_cnnp_model_t* const super, ccv_nnc_symbolic_graph_t* const graph, const ccv_nnc_tensor_symbol_t* const inputs, const int input_size, ccv_nnc_tensor_symbol_t* const outputs, const int output_size)

{

  assert(input_size == 1);

  assert(output_size == 1);

  ccv_cnnp_model_permute_t* const self = (ccv_cnnp_model_permute_t*)super;

  PRINT(CCV_CLI_VERBOSE, "[cnnp_permute_build] -\n");

  ccv_nnc_tensor_param_t params = ccv_nnc_tensor_symbol_params(graph, inputs[0]);

  ccv_nnc_tensor_symbol_t to = ccv_nnc_tensor_symbol_alias_to(graph, inputs[0]);

  const int nd = ccv_nnc_tensor_nd(params.dim);

  int input_dim[CCV_NNC_MAX_DIM_ALLOC];

  memcpy(input_dim, params.dim, sizeof(params.dim));

  int input_stride[CCV_NNC_MAX_DIM_ALLOC] = {};

  int output_stride[CCV_NNC_MAX_DIM_ALLOC] = {};

  if (to.d == CCV_NNC_NO_TENSOR_SYMBOL) // If it is not an alias. Find stride and permute.

  {

    ccv_nnc_tensor_get_stride(input_dim, input_stride);

    int i;

    for (i = 0; i < nd; i++)

    {

      const int idx = self->index[i];

      assert(idx >= 0 && idx < nd);

      params.dim[i] = input_dim[idx];

      output_stride[i] = input_stride[idx];

    }

    outputs[0] = ccv_nnc_tensor_symbol_alias_new(graph, inputs[0], ccv_nnc_no_ofs, output_stride, params, 0);

  } else {

    // if it is an alias, we can get the stride from it and use that.

    int input_ofs[CCV_NNC_MAX_DIM_ALLOC];

    ccv_nnc_tensor_symbol_alias_params(graph, inputs[0], input_ofs, input_stride);

    assert(input_stride[0] != 0);

    int output_ofs[CCV_NNC_MAX_DIM_ALLOC] = {};

    int i;

    for (i = 0; i < nd; 
i++3
)

    {

      const int idx = self->index[i];

      assert(idx >= 0 && idx < nd);

      params.dim[i] = input_dim[idx];

      output_stride[i] = input_stride[idx];

      output_ofs[i] = input_ofs[idx];

    }

    outputs[0] = ccv_nnc_tensor_symbol_alias_new(graph, inputs[0], output_ofs, output_stride, params, 0);

  }

}

static ccv_cnnp_model_t* _ccv_cnnp_permute_copy(const ccv_cnnp_model_t* const super, void* const context);

static const ccv_cnnp_model_vtab_t ccv_cnnp_permute_isa = {

  .build = _ccv_cnnp_permute_build,

  .copy = _ccv_cnnp_permute_copy,

};

ccv_cnnp_model_t* ccv_cnnp_permute(const int index[CCV_NNC_MAX_DIM_ALLOC], const char* const name)

{

  ccv_cnnp_model_permute_t* const model_permute = (ccv_cnnp_model_permute_t*)cccalloc(1, sizeof(ccv_cnnp_model_permute_t));

  model_permute->super.isa = &ccv_cnnp_permute_isa;

  model_permute->super.input_size = 1;

  model_permute->super.outputs = &model_permute->output;

  model_permute->super.output_size = 1;

  ccv_cnnp_model_copy_name(&model_permute->super, name);

  memcpy(model_permute->index, index, sizeof(model_permute->index));

  return (ccv_cnnp_model_t*)model_permute;

}

static ccv_cnnp_model_t* _ccv_cnnp_permute_copy(const ccv_cnnp_model_t* const super, void* const context)

{

  const ccv_cnnp_model_permute_t* const self = (const ccv_cnnp_model_permute_t*)super;

  return ccv_cnnp_permute(self->index, self->super.name);

}

typedef struct {

  ccv_cnnp_model_t super;

  int index;

  ccv_nnc_tensor_symbol_t output;

} ccv_cnnp_model_extract_t;

static void _ccv_cnnp_extract_build(ccv_cnnp_model_t* const super, ccv_nnc_symbolic_graph_t* const graph, const ccv_nnc_tensor_symbol_t* const inputs, const int input_size, ccv_nnc_tensor_symbol_t* const outputs, const int output_size)

{

  assert(output_size == 1);

  ccv_cnnp_model_extract_t* const self = (ccv_cnnp_model_extract_t*)super;

  PRINT(CCV_CLI_VERBOSE, "[cnnp_extract_build] index: %d\n", self->index);

  outputs[0] = inputs[self->index];

}

static ccv_cnnp_model_t* _ccv_cnnp_extract_copy(const ccv_cnnp_model_t* const self, void* const context);

static const ccv_cnnp_model_vtab_t ccv_cnnp_extract_isa = {

  .build = _ccv_cnnp_extract_build,

  .copy = _ccv_cnnp_extract_copy,

};

ccv_cnnp_model_t* ccv_cnnp_extract(const int index, const char* const name)

{

  ccv_cnnp_model_extract_t* const model_extract = (ccv_cnnp_model_extract_t*)cccalloc(1, sizeof(ccv_cnnp_model_extract_t));

  model_extract->index = index;

  model_extract->super.isa = &ccv_cnnp_extract_isa;

  model_extract->super.input_size = 0;

  model_extract->super.outputs = &model_extract->output;

  model_extract->super.output_size = 1;

  ccv_cnnp_model_copy_name(&model_extract->super, name);

  return (ccv_cnnp_model_t*)model_extract;

}

static ccv_cnnp_model_t* _ccv_cnnp_extract_copy(const ccv_cnnp_model_t* const super, void* const context)

{

  ccv_cnnp_model_extract_t* const self = (ccv_cnnp_model_extract_t*)super;

  return ccv_cnnp_extract(self->index, self->super.name);

}

typedef struct {

  ccv_cnnp_model_t super;

  ccv_nnc_tensor_symbol_t output;

} ccv_cnnp_model_flatten_t;

static void _ccv_cnnp_flatten_build(ccv_cnnp_model_t* const super, ccv_nnc_symbolic_graph_t* const graph, const ccv_nnc_tensor_symbol_t* const inputs, const int input_size, ccv_nnc_tensor_symbol_t* const outputs, const int output_size)

{

  PRINT(CCV_CLI_VERBOSE, "[cnnp_flatten_build] -\n");

  assert(input_size == 1);

  assert(output_size == 1);

  const ccv_nnc_tensor_param_t params = ccv_nnc_tensor_symbol_params(graph, inputs[0]);

  ccv_nnc_tensor_param_t output_params = params;

  memset(output_params.dim, 0, sizeof(output_params.dim));

  output_params.dim[0] = ccv_nnc_tensor_get_n(params);

  assert(output_params.dim[0] > 0);

  output_params.dim[1] = ccv_nnc_tensor_count(params) / output_params.dim[0];

  int stride[CCV_NNC_MAX_DIM_ALLOC] = {};