Function TEST

tensorflow/lite/delegates/gpu/gl/kernels/softmax_test.cc:34–53

namespace {

TEST(SoftmaxTest, Softmax) {
  TensorRef<BHWC> input;
  input.type = DataType::FLOAT32;
  input.ref = 0;
  input.shape = BHWC(1, 2, 2, 1);

  TensorRef<BHWC> output;
  output.type = DataType::FLOAT32;
  output.ref = 1;
  output.shape = BHWC(1, 2, 2, 1);

  SoftmaxAttributes attr;
  attr.axis = Axis::CHANNELS;

  SingleOpModel model({ToString(OperationType::SOFTMAX), attr}, {input},
                      {output});
  ASSERT_TRUE(model.PopulateTensor(0, {0.1, 0.2, 0.1, 0.2}));
  ASSERT_OK(model.Invoke(*NewSoftmaxNodeShader()));
  EXPECT_THAT(model.GetOutput(0), Pointwise(FloatNear(1e-6), {1, 1, 1, 1}));
}

TEST(SoftmaxTest, DoesNotWorkForHeightAxis) {
  TensorRef<BHWC> input;