Repository: microsoft/onnxruntime
File: onnxruntime/core/providers/coreml/builders/impl/gather_op_builder.cc (C++)

Code before:
    const auto validate_indices = false;
    AddOperationInput(*op, "x", node.InputDefs()[0]->Name());                                   // data
    AddOperationInput(*op, "indices", node.InputDefs()[1]->Name());                             // indices
    AddOperationInput(*op, "axis", model_builder.AddScalarConstant(op->type(), "axis", axis));  // axis attr
    AddOperationInput(*op, "validate_indices", model_builder.AddScalarConstant(op->type(), "validate_indices", validate_indices));
    AddOperationOutput(*op, *node.OutputDefs()[0]);  // output
    model_builder.AddOperation(std::move(op));
  } else {
    auto layer = model_builder.CreateNNLayer(node);
    layer->mutable_gather()->set_axis(GetAxisAttribute(node));
    *layer->mutable_input()->Add() = node.InputDefs()[0]->Name();    // data
    *layer->mutable_input()->Add() = node.InputDefs()[1]->Name();    // indices
    *layer->mutable_output()->Add() = node.OutputDefs()[0]->Name();  // output
    model_builder.AddLayer(std::move(layer));
  }
  return Status::OK();
}

bool GatherOpBuilder::HasSupportedInputsImpl(const Node& node, const OpBuilderInputParams& /*input_params*/,
                                             const logging::Logger& logger) const {
  int32_t input_type;
  if (!GetType(*node.InputDefs()[0], input_type, logger))
    return false;

  if (input_type != ONNX_NAMESPACE::TensorProto_DataType_FLOAT &&
      input_type != ONNX_NAMESPACE::TensorProto_DataType_FLOAT16 &&
      input_type != ONNX_NAMESPACE::TensorProto_DataType_INT64) {
    LOGS(logger, VERBOSE) << "[" << node.OpType()
                          << "] Input type: [" << input_type
                          << "] is not supported for now";
    return false;
  }

  return true;
}

bool GatherOpBuilder::IsOpSupportedImpl(const Node& node, const OpBuilderInputParams& /*input_params*/,
                                        const logging::Logger& logger) const {
  std::vector<int64_t> data_shape, indices_shape;
  if (!GetShape(*node.InputDefs()[0], data_shape, logger)) {
    LOGS(logger, VERBOSE) << "Failed to get 'data' shape";
    return false;
  }

  if (!GetShape(*node.InputDefs()[1], indices_shape, logger)) {
    LOGS(logger, VERBOSE) << "Failed to get 'indices' shape";
    return false;
  }

  // Don't allow scalar 'indices' input.
  // We convert scalar inputs to tensors with shape [1] before providing them to CoreML.

Review comment (suggestion):
should allowing float16 should be limited to the ML program case? that info is available in `OpBuilderInputParams::create_mlprogram`.



Code after:
    const auto validate_indices = false;
    AddOperationInput(*op, "x", node.InputDefs()[0]->Name());                                   // data
    AddOperationInput(*op, "indices", node.InputDefs()[1]->Name());                             // indices
    AddOperationInput(*op, "axis", model_builder.AddScalarConstant(op->type(), "axis", axis));  // axis attr
    AddOperationInput(*op, "validate_indices", model_builder.AddScalarConstant(op->type(), "validate_indices", validate_indices));
    AddOperationOutput(*op, *node.OutputDefs()[0]);  // output
    model_builder.AddOperation(std::move(op));
  } else {
    auto layer = model_builder.CreateNNLayer(node);
    layer->mutable_gather()->set_axis(GetAxisAttribute(node));
    *layer->mutable_input()->Add() = node.InputDefs()[0]->Name();    // data
    *layer->mutable_input()->Add() = node.InputDefs()[1]->Name();    // indices
    *layer->mutable_output()->Add() = node.OutputDefs()[0]->Name();  // output
    model_builder.AddLayer(std::move(layer));
  }
  return Status::OK();
}

bool GatherOpBuilder::HasSupportedInputsImpl(const Node& node, const OpBuilderInputParams& input_params,
                                             const logging::Logger& logger) const {
  int32_t input_type;
  if (!GetType(*node.InputDefs()[0], input_type, logger))
    return false;

  if (input_type != ONNX_NAMESPACE::TensorProto_DataType_FLOAT &&
      (input_type != ONNX_NAMESPACE::TensorProto_DataType_FLOAT16 ||
       !input_params.create_mlprogram || input_params.coreml_version < 6) &&
      input_type != ONNX_NAMESPACE::TensorProto_DataType_INT64) {
    LOGS(logger, VERBOSE) << "[" << node.OpType()
                          << "] Input type: [" << input_type
                          << "] is not supported for now";
    return false;
  }

  return true;
}

bool GatherOpBuilder::IsOpSupportedImpl(const Node& node, const OpBuilderInputParams& /*input_params*/,
                                        const logging::Logger& logger) const {
  std::vector<int64_t> data_shape, indices_shape;
  if (!GetShape(*node.InputDefs()[0], data_shape, logger)) {
    LOGS(logger, VERBOSE) << "Failed to get 'data' shape";
    return false;
  }

  if (!GetShape(*node.InputDefs()[1], indices_shape, logger)) {
    LOGS(logger, VERBOSE) << "Failed to get 'indices' shape";
    return false;
  }

  // Don't allow scalar 'indices' input.