ctx, vals.data(), vals_shape, 1, &sum_dims_raw));
    sum_dims.reset(sum_dims_raw);
  }

  // Reduce sum the output on all dimensions.
  TF_RETURN_IF_ERROR(ops::Sum(ctx, model_out.get(), sum_dims.get(), &outputs[0],
                              /*keep_dims=*/false, "sum_output"));
  return absl::OkStatus();
}
// ========================= End Helper Functions==============================