#include <ArgMax.h>

Collaboration diagram for luci_interpreter::kernels::ArgMax:

[legend]

Public Member Functions
	ArgMax (const Tensor input, const Tensor axis, Tensor *output, const ArgMaxParams &params)

const Tensor *	input () const

const Tensor *	axis () const

Tensor *	output () const

void	configure () override

void	execute () const override

Public Member Functions inherited from luci_interpreter::KernelWithParams< ArgMaxParams >
const ArgMaxParams &	params () const

Public Member Functions inherited from luci_interpreter::Kernel
virtual	~Kernel ()=default

const std::vector< const Tensor * > &	getInputTensors () const

const std::vector< Tensor * > &	getOutputTensors () const

Additional Inherited Members
Protected Member Functions inherited from luci_interpreter::KernelWithParams< ArgMaxParams >
	KernelWithParams (std::vector< const Tensor * > inputs, std::vector< Tensor * > outputs, const ArgMaxParams &params)

Protected Member Functions inherited from luci_interpreter::Kernel
	Kernel (std::vector< const Tensor * > inputs, std::vector< Tensor * > outputs)

Protected Attributes inherited from luci_interpreter::KernelWithParams< ArgMaxParams >
const ArgMaxParams	_params

Protected Attributes inherited from luci_interpreter::Kernel
const std::vector< const Tensor * >	_inputs

const std::vector< Tensor * >	_outputs

Detailed Description

Definition at line 28 of file ArgMax.h.

Constructor & Destructor Documentation

◆ ArgMax()

luci_interpreter::kernels::ArgMax::ArgMax	(	const Tensor *	input,
		const Tensor *	axis,
		Tensor *	output,
		const ArgMaxParams &	params
	)

Definition at line 26 of file ArgMax.cpp.

  : KernelWithParams<ArgMaxParams>({input, axis}, {output}, params)
{
}

References axis(), and input().

Member Function Documentation

◆ axis()

const Tensor * luci_interpreter::kernels::ArgMax::axis ( ) const

inline

Definition at line 34 of file ArgMax.h.

34{ return _inputs[1]; }

luci_interpreter::Kernel::_inputs

const std::vector< const Tensor * > _inputs

Definition Kernel.h:52

References luci_interpreter::Kernel::_inputs.

Referenced by ArgMax(), configure(), and execute().

◆ configure()

void luci_interpreter::kernels::ArgMax::configure ( )

overridevirtual

Implements luci_interpreter::Kernel.

Definition at line 31 of file ArgMax.cpp.

{
  assert(axis()->element_type() == DataType::S32 || axis()->element_type() == DataType::S64);
  assert(input()->shape().num_dims() >= 1);
  const Shape &input_shape = input()->shape();
  const int num_dims = input_shape.num_dims();
  Shape output_shape(num_dims - 1);
 
  // If axis value is negative, then update by adding input_shape's num_dims.
  // If updated value also negative, then assert.
  assert(axis()->shape().num_elements() == 1);
  int axis_value = getTensorData<int32_t>(axis())[0];
  if (axis_value < 0)
    axis_value = axis_value + num_dims;
  assert(axis_value >= 0);
 
  int j = 0;
  for (int i = 0; i < num_dims; i++)
  {
    if (i == axis_value)
      continue;
    output_shape.dim(j++) = input_shape.dim(i);
  }
 
  assert(output()->element_type() == _params.output_type);
 
  output()->resize(output_shape);
}

References luci_interpreter::KernelWithParams< ArgMaxParams >::_params, axis(), luci_interpreter::Shape::dim(), input(), luci_interpreter::Shape::num_dims(), output(), output_shape, luci_interpreter::ArgMaxParams::output_type, luci_interpreter::Tensor::resize(), and luci_interpreter::Tensor::shape().

◆ execute()

void luci_interpreter::kernels::ArgMax::execute ( ) const

overridevirtual

Implements luci_interpreter::Kernel.

Definition at line 60 of file ArgMax.cpp.

{
 
#define TF_LITE_ARG_MAX(data_type, axis_type, output_type)                                    \
  luci_interpreter_pal::ArgMinMax(getTensorShape(input()), getTensorData<data_type>(input()), \
                                  getTensorData<axis_type>(axis()), getTensorShape(output()), \
                                  getTensorData<output_type>(output()), std::greater<data_type>())
  if (axis()->element_type() == DataType::S32)
  {
    switch (_params.output_type)
    {
      case DataType::S32:
        switch (input()->element_type())
        {
          case DataType::FLOAT32:
            TF_LITE_ARG_MAX(float, int32_t, int32_t);
            break;
          case DataType::U8:
            TF_LITE_ARG_MAX(uint8_t, int32_t, int32_t);
            break;
          default:
            throw std::runtime_error("Unsupported input type.");
        }
        break;
      case DataType::S64:
        switch (input()->element_type())
        {
          case DataType::FLOAT32:
            TF_LITE_ARG_MAX(float, int32_t, int64_t);
            break;
          case DataType::U8:
            TF_LITE_ARG_MAX(uint8_t, int32_t, int64_t);
            break;
          default:
            throw std::runtime_error("Unsupported input type.");
        }
        break;
      default:
        throw std::runtime_error("Unsupported output type.");
    }
  }
  else
  {
    switch (_params.output_type)
    {
      case DataType::S32:
        switch (input()->element_type())
        {
          case DataType::FLOAT32:
            TF_LITE_ARG_MAX(float, int64_t, int32_t);
            break;
          case DataType::U8:
            TF_LITE_ARG_MAX(uint8_t, int64_t, int32_t);
            break;
          default:
            throw std::runtime_error("Unsupported input type.");
        }
        break;
      case DataType::S64:
        switch (input()->element_type())
        {
          case DataType::FLOAT32:
            TF_LITE_ARG_MAX(float, int64_t, int64_t);
            break;
          case DataType::U8:
            TF_LITE_ARG_MAX(uint8_t, int64_t, int64_t);
            break;
          default:
            throw std::runtime_error("Unsupported input type.");
        }
        break;
      default:
        throw std::runtime_error("Unsupported output type.");
    }
  }
#undef TF_LITE_ARG_MAX
}