/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/core/framework/common_shape_fns.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/shape_inference.h"

namespace tensorflow {
namespace {

// Verifies that `shapes_and_types` is a valid list handle and has the right
// dtype.
Status VerifyHandleData(
    shape_inference::InferenceContext* c,
    const std::vector<shape_inference::ShapeAndType>& shapes_and_types,
    DataType element_dtype) {
  if (shapes_and_types.size() != 1) {
    return errors::InvalidArgument(
        "Invalid handle_data for input list. Expected length of "
        "shape_and_types: ",
        1, " Saw: ", shapes_and_types.size());
  }
  const shape_inference::ShapeAndType& list_shape_type = shapes_and_types[0];
  if (list_shape_type.dtype != element_dtype) {
    return errors::InvalidArgument("Expected list with element dtype ",
                                   DataTypeString(element_dtype),
                                   " but got list with element dtype ",
                                   DataTypeString(list_shape_type.dtype));
  }
  return Status::OK();
}

// Assumes that the handle_data is valid.
shape_inference::ShapeHandle GetElementShapeFromHandleData(
    const std::vector<shape_inference::ShapeAndType>& shapes_and_types) {
  return shapes_and_types[0].shape;
}

REGISTER_OP("EmptyTensorList")
    .Input("element_shape: shape_type")
    .Input("max_num_elements: int32")
    .Output("handle: variant")
    .Attr("element_dtype: type")
    .Attr("shape_type: {int32, int64}")
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      c->set_output(0, c->Scalar());
      DataType element_dtype;
      TF_RETURN_IF_ERROR(c->GetAttr("element_dtype", &element_dtype));
      shape_inference::ShapeHandle element_shape;
      TF_RETURN_IF_ERROR(c->MakeShapeFromShapeTensorTreatScalarAsUnknownShape(
          0, &element_shape));
      c->set_output_handle_shapes_and_types(
          0, std::vector<shape_inference::ShapeAndType>{
                 {element_shape, element_dtype}});
      return Status::OK();
    });

REGISTER_OP("TensorListPushBack")
    .Input("input_handle: variant")
    .Input("tensor: element_dtype")
    .Output("output_handle: variant")
    .Attr("element_dtype: type")
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      c->set_output(0, c->Scalar());
      DataType element_dtype;
      TF_RETURN_IF_ERROR(c->GetAttr("element_dtype", &element_dtype));
      shape_inference::ShapeHandle element_shape = c->UnknownShape();

      auto* handle_data = c->input_handle_shapes_and_types(0);
      if (handle_data != nullptr && handle_data->size() > 1) {
        return errors::InvalidArgument(
            "Trying to push to list with wrong variant data.");
      }
      if (handle_data != nullptr && handle_data->size() == 1) {
        const shape_inference::ShapeAndType& list_shape_type =
            (*handle_data)[0];
        if (list_shape_type.dtype != element_dtype) {
          return errors::InvalidArgument(
              "Trying to push to list with wrong element dtype. List has type ",
              DataTypeString(list_shape_type.dtype),
              " but trying to push element with type ",
              DataTypeString(element_dtype));
        }
        shape_inference::ShapeHandle ignored;
        TF_RETURN_IF_ERROR(
            c->Merge(element_shape, list_shape_type.shape, &ignored));
        element_shape = list_shape_type.shape;
      }
      c->set_output_handle_shapes_and_types(
          0, std::vector<shape_inference::ShapeAndType>{
                 {element_shape, element_dtype}});
      return Status::OK();
    });

REGISTER_OP("TensorListPushBackBatch")
    .Input("input_handles: variant")
    .Input("tensor: element_dtype")
    .Output("output_handles: variant")
    .Attr("element_dtype: type")
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      shape_inference::ShapeHandle input_handles;
      TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 1, &input_handles));

      shape_inference::ShapeHandle tensor;
      TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(1), 1, &tensor));

      TF_RETURN_IF_ERROR(
          c->MergePrefix(tensor, input_handles, &tensor, &input_handles));

      c->set_output(0, input_handles);

      DataType element_dtype;
      TF_RETURN_IF_ERROR(c->GetAttr("element_dtype", &element_dtype));
      shape_inference::ShapeHandle element_shape = c->UnknownShape();

      auto* handle_data = c->input_handle_shapes_and_types(0);
      if (handle_data != nullptr && handle_data->size() > 1) {
        return errors::InvalidArgument(
            "Trying to push to list with wrong variant data.");
      }
      if (handle_data != nullptr && handle_data->size() == 1) {
        const shape_inference::ShapeAndType& list_shape_type =
            (*handle_data)[0];
        if (list_shape_type.dtype != element_dtype) {
          return errors::InvalidArgument(
              "Trying to push to list with wrong element dtype. List has type ",
              DataTypeString(list_shape_type.dtype),
              " but trying to push element with type ",
              DataTypeString(element_dtype));
        }
        shape_inference::ShapeHandle ignored;
        TF_RETURN_IF_ERROR(
            c->Merge(element_shape, list_shape_type.shape, &ignored));
        element_shape = list_shape_type.shape;
      }
      c->set_output_handle_shapes_and_types(
          0, std::vector<shape_inference::ShapeAndType>{
                 {element_shape, element_dtype}});
      return Status::OK();
    });

REGISTER_OP("TensorListLength")
    .Input("input_handle: variant")
    .Output("length: int32")
    .SetShapeFn(shape_inference::ScalarShape);

REGISTER_OP("TensorListPopBack")
    .Input("input_handle: variant")
    .Input("element_shape: int32")
    .Output("output_handle: variant")
    .Output("tensor: element_dtype")
    .Attr("element_dtype: type")
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      DataType element_dtype;
      TF_RETURN_IF_ERROR(c->GetAttr("element_dtype", &element_dtype));
      shape_inference::ShapeHandle tensor_shape = c->UnknownShape();
      auto* handle_data = c->input_handle_shapes_and_types(0);
      if (handle_data != nullptr && handle_data->size() > 1) {
        return errors::InvalidArgument(
            "Trying to read from list with invalid variant data.");
      }
      if (handle_data != nullptr && handle_data->size() == 1) {
        const shape_inference::ShapeAndType& list_shape_type =
            (*handle_data)[0];
        if (list_shape_type.dtype != element_dtype) {
          return errors::InvalidArgument(
              "Trying to read from list with wrong element dtype. List has "
              "type ",
              DataTypeString(list_shape_type.dtype),
              " but trying to push element with type ",
              DataTypeString(element_dtype));
        }
        shape_inference::ShapeHandle ignored;
        TF_RETURN_IF_ERROR(
            c->Merge(tensor_shape, list_shape_type.shape, &ignored));
        c->set_output_handle_shapes_and_types(0, *handle_data);
        tensor_shape = list_shape_type.shape;
      }
      c->set_output(1, tensor_shape);
      c->set_output(0, c->Scalar());
      return Status::OK();
    });

REGISTER_OP("TensorListStack")
    .Input("input_handle: variant")
    .Input("element_shape: int32")
    .Output("tensor: element_dtype")
    .Attr("element_dtype: type")
    .Attr("num_elements: int = -1")
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      DataType element_dtype;
      TF_RETURN_IF_ERROR(c->GetAttr("element_dtype", &element_dtype));
      shape_inference::ShapeHandle element_shape = c->UnknownShape();
      auto* handle_data = c->input_handle_shapes_and_types(0);
      if (handle_data != nullptr && handle_data->size() > 1) {
        return errors::InvalidArgument(
            "Trying to read from list with wrong variant data.");
      }
      if (handle_data != nullptr && handle_data->size() == 1) {
        const shape_inference::ShapeAndType& list_shape_type =
            (*handle_data)[0];
        if (list_shape_type.dtype != element_dtype) {
          return errors::InvalidArgument(
              "Trying to read from list with wrong element dtype. List has "
              "type ",
              DataTypeString(list_shape_type.dtype), " but expected type ",
              DataTypeString(element_dtype));
        }
        shape_inference::ShapeHandle ignored;
        TF_RETURN_IF_ERROR(
            c->Merge(element_shape, list_shape_type.shape, &ignored));
        element_shape = list_shape_type.shape;
      }
      shape_inference::ShapeHandle element_shape_input = c->UnknownShape();
      TF_RETURN_IF_ERROR(c->MakeShapeFromShapeTensorTreatScalarAsUnknownShape(
          1, &element_shape_input));
      TF_RETURN_IF_ERROR(
          c->Merge(element_shape, element_shape_input, &element_shape));
      int expected_num_elements = -1;
      TF_RETURN_IF_ERROR(c->GetAttr("num_elements", &expected_num_elements));
      shape_inference::ShapeHandle num_elements;
      if (expected_num_elements == -1) {
        num_elements = c->MakeShape({c->UnknownDim()});
      } else {
        num_elements = c->MakeShape({expected_num_elements});
      }
      shape_inference::ShapeHandle result;
      TF_RETURN_IF_ERROR(c->Concatenate(num_elements, element_shape, &result));
      c->set_output(0, result);
      return Status::OK();
    });

Status TensorListConcatShapeInference(
    shape_inference::InferenceContext* c,
    shape_inference::ShapeHandle element_shape) {
  DataType element_dtype;
  TF_RETURN_IF_ERROR(c->GetAttr("element_dtype", &element_dtype));
  auto* handle_data = c->input_handle_shapes_and_types(0);
  if (handle_data != nullptr && handle_data->size() > 1) {
    return errors::InvalidArgument(
        "Trying to read from list with wrong variant data.");
  }
  if (handle_data != nullptr && handle_data->size() == 1) {
    const shape_inference::ShapeAndType& list_shape_type = (*handle_data)[0];
    if (list_shape_type.dtype != element_dtype) {
      return errors::InvalidArgument(
          "Trying to read from list with wrong element dtype. List has "
          "type ",
          DataTypeString(list_shape_type.dtype), " but expected type ",
          DataTypeString(element_dtype));
    }
    shape_inference::ShapeHandle merged;
    TF_RETURN_IF_ERROR(c->Merge(element_shape, list_shape_type.shape, &merged));
    element_shape = merged;
  }
  if (c->RankKnown(element_shape)) {
    shape_inference::ShapeHandle result;
    TF_RETURN_IF_ERROR(c->Subshape(element_shape, 1, &result));
    TF_RETURN_IF_ERROR(
        c->Concatenate(c->MakeShape({c->UnknownDim()}), result, &result));
    c->set_output(0, result);
  } else {
    c->set_output(0, c->UnknownShape());
  }
  c->set_output(1, c->MakeShape({c->UnknownDim()}));
  return Status::OK();
}

REGISTER_OP("TensorListConcat")
    .Input("input_handle: variant")
    .Output("tensor: element_dtype")
    .Output("lengths: int64")
    .Attr("element_dtype: type")
    .Attr("element_shape: shape = { unknown_rank: true }")
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      PartialTensorShape raw_element_shape;
      TF_RETURN_IF_ERROR(c->GetAttr("element_shape", &raw_element_shape));
      shape_inference::ShapeHandle element_shape;
      TF_RETURN_IF_ERROR(c->MakeShapeFromPartialTensorShape(raw_element_shape,
                                                            &element_shape));
      return TensorListConcatShapeInference(c, element_shape);
    });

REGISTER_OP("TensorListConcatV2")
    .Input("input_handle: variant")
    .Input("element_shape: shape_type")
    .Input("leading_dims: int64")
    .Output("tensor: element_dtype")
    .Output("lengths: int64")
    .Attr("element_dtype: type")
    .Attr("shape_type: {int32, int64}")
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      shape_inference::ShapeHandle element_shape;
      TF_RETURN_IF_ERROR(c->MakeShapeFromShapeTensorTreatScalarAsUnknownShape(
          1, &element_shape));
      return TensorListConcatShapeInference(c, element_shape);
    });

REGISTER_OP("TensorListSplit")
    .Input("tensor: element_dtype")
    .Input("element_shape: shape_type")
    .Input("lengths: int64")
    .Output("output_handle: variant")
    .Attr("element_dtype: type")
    .Attr("shape_type: {int32, int64}")
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      c->set_output(0, c->Scalar());
      DataType element_dtype;
      TF_RETURN_IF_ERROR(c->GetAttr("element_dtype", &element_dtype));
      shape_inference::ShapeHandle tensor_shape = c->input(0);
      shape_inference::ShapeHandle ignored;
      // Check that tensor is at least a vector.
      TF_RETURN_IF_ERROR(c->WithRankAtLeast(tensor_shape, 1, &ignored));
      // Check that lengths is a vector.
      TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 1, &ignored));
      shape_inference::ShapeHandle element_shape_from_tensor_shape;
      TF_RETURN_IF_ERROR(
          c->Subshape(tensor_shape, 1, &element_shape_from_tensor_shape));
      TF_RETURN_IF_ERROR(c->Concatenate(c->MakeShape({c->UnknownDim()}),
                                        element_shape_from_tensor_shape,
                                        &element_shape_from_tensor_shape));
      shape_inference::ShapeHandle element_shape;
      TF_RETURN_IF_ERROR(c->MakeShapeFromShapeTensorTreatScalarAsUnknownShape(
          1, &element_shape));
      TF_RETURN_IF_ERROR(c->Merge(element_shape_from_tensor_shape,
                                  element_shape,
                                  &element_shape_from_tensor_shape));
      c->set_output_handle_shapes_and_types(
          0, std::vector<shape_inference::ShapeAndType>{
                 {element_shape, element_dtype}});
      return Status::OK();
    });

REGISTER_OP("TensorListFromTensor")
    .Input("tensor: element_dtype")
    .Input("element_shape: shape_type")
    .Output("output_handle: variant")
    .Attr("element_dtype: type")
    .Attr("shape_type: {int32, int64}")
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      c->set_output(0, c->Scalar());
      DataType element_dtype;
      TF_RETURN_IF_ERROR(c->GetAttr("element_dtype", &element_dtype));
      shape_inference::ShapeHandle tensor_shape = c->input(0);
      shape_inference::ShapeHandle tensor_shape_except_first_dim;
      TF_RETURN_IF_ERROR(
          c->Subshape(tensor_shape, 1, &tensor_shape_except_first_dim));
      shape_inference::ShapeHandle element_shape;
      TF_RETURN_IF_ERROR(c->MakeShapeFromShapeTensorTreatScalarAsUnknownShape(
          1, &element_shape));
      TF_RETURN_IF_ERROR(c->Merge(tensor_shape_except_first_dim, element_shape,
                                  &tensor_shape_except_first_dim));
      c->set_output_handle_shapes_and_types(
          0, std::vector<shape_inference::ShapeAndType>{
                 {element_shape, element_dtype}});
      return Status::OK();
    });

REGISTER_OP("TensorListElementShape")
    .Input("input_handle: variant")
    .Output("element_shape: shape_type")
    .Attr("shape_type: {int32, int64}")
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      auto* handle_data = c->input_handle_shapes_and_types(0);
      if (handle_data == nullptr) {
        c->set_output(0, c->Vector(c->UnknownDim()));
        return Status::OK();
      }
      c->set_output(0, c->Vector(c->Rank((*handle_data)[0].shape)));
      return Status::OK();
    });

REGISTER_OP("TensorListReserve")
    .Input("element_shape: shape_type")
    .Input("num_elements: int32")
    .Output("handle: variant")
    .Attr("element_dtype: type")
    .Attr("shape_type: {int32, int64}")
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      c->set_output(0, c->Scalar());
      shape_inference::ShapeHandle element_shape;
      TF_RETURN_IF_ERROR(c->MakeShapeFromShapeTensorTreatScalarAsUnknownShape(
          0, &element_shape));
      DataType element_dtype;
      TF_RETURN_IF_ERROR(c->GetAttr("element_dtype", &element_dtype));
      c->set_output_handle_shapes_and_types(
          0, std::vector<shape_inference::ShapeAndType>{
                 {element_shape, element_dtype}});
      return Status::OK();
    });

REGISTER_OP("TensorListGetItem")
    .Input("input_handle: variant")
    .Input("index: int32")
    .Input("element_shape: int32")
    .Output("item: element_dtype")
    .Attr("element_dtype: type")
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      DataType element_dtype;
      TF_RETURN_IF_ERROR(c->GetAttr("element_dtype", &element_dtype));
      auto* handle_data = c->input_handle_shapes_and_types(0);
      shape_inference::ShapeHandle element_shape = c->UnknownShape();
      if (handle_data != nullptr) {
        const shape_inference::ShapeAndType& list_shape_type =
            (*handle_data)[0];
        element_shape = list_shape_type.shape;
        if (list_shape_type.dtype != element_dtype) {
          return errors::InvalidArgument("Expected list with element dtype ",
                                         DataTypeString(element_dtype),
                                         " but got list with element dtype ",
                                         DataTypeString(list_shape_type.dtype));
        }
      }
      shape_inference::ShapeHandle element_shape_input = c->UnknownShape();
      TF_RETURN_IF_ERROR(c->MakeShapeFromShapeTensorTreatScalarAsUnknownShape(
          2, &element_shape_input));
      TF_RETURN_IF_ERROR(
          c->Merge(element_shape, element_shape_input, &element_shape));
      c->set_output(0, element_shape);
      return Status::OK();
    });

REGISTER_OP("TensorListResize")
    .Input("input_handle: variant")
    .Input("size: int32")
    .Output("output_handle: variant")
    .SetIsStateful()
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      // Check that `size` has scalar shape.
      shape_inference::ShapeHandle size_shape = c->input(1);
      shape_inference::ShapeHandle unused;
      TF_RETURN_IF_ERROR(c->WithRank(size_shape, 0, &unused));
      c->set_output(0, c->Scalar());
      auto* handle_data = c->input_handle_shapes_and_types(0);
      if (handle_data != nullptr) {
        c->set_output_handle_shapes_and_types(0, *handle_data);
      }
      return Status::OK();
    });

REGISTER_OP("TensorListSetItem")
    .Input("input_handle: variant")
    .Input("index: int32")
    .Input("item: element_dtype")
    .Output("output_handle: variant")
    .Attr("element_dtype: type")
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      DataType element_dtype;
      TF_RETURN_IF_ERROR(c->GetAttr("element_dtype", &element_dtype));
      auto* handle_data = c->input_handle_shapes_and_types(0);
      c->set_output(0, c->Scalar());
      if (handle_data == nullptr) {
        c->set_output_handle_shapes_and_types(
            0, {{c->UnknownShape(), element_dtype}});
        return Status::OK();
      }
      const shape_inference::ShapeAndType& list_shape_type = (*handle_data)[0];
      shape_inference::ShapeHandle item_shape = c->input(2);
      TF_RETURN_IF_ERROR(
          c->Merge(item_shape, list_shape_type.shape, &item_shape));
      c->set_output_handle_shapes_and_types(0, *handle_data);
      return Status::OK();
    });

REGISTER_OP("TensorListGather")
    .Input("input_handle: variant")
    .Input("indices: int32")
    .Input("element_shape: int32")
    .Output("values: element_dtype")
    .Attr("element_dtype: type")
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      DataType element_dtype;
      TF_RETURN_IF_ERROR(c->GetAttr("element_dtype", &element_dtype));
      auto* handle_data = c->input_handle_shapes_and_types(0);
      shape_inference::ShapeHandle element_shape = c->UnknownShape();
      if (handle_data != nullptr) {
        const shape_inference::ShapeAndType& list_shape_type =
            (*handle_data)[0];
        element_shape = list_shape_type.shape;
        if (list_shape_type.dtype != element_dtype) {
          return errors::InvalidArgument("Expected list with element dtype ",
                                         DataTypeString(element_dtype),
                                         " but got list with element dtype ",
                                         DataTypeString(list_shape_type.dtype));
        }
      }
      shape_inference::ShapeHandle element_shape_input = c->UnknownShape();
      TF_RETURN_IF_ERROR(c->MakeShapeFromShapeTensorTreatScalarAsUnknownShape(
          2, &element_shape_input));
      TF_RETURN_IF_ERROR(
          c->Merge(element_shape, element_shape_input, &element_shape));
      shape_inference::ShapeHandle out;
      TF_RETURN_IF_ERROR(c->Concatenate(c->input(1), element_shape, &out));
      c->set_output(0, out);
      return Status::OK();
    });

REGISTER_OP("TensorListScatter")
    .Input("tensor: element_dtype")
    .Input("indices: int32")
    .Input("element_shape: shape_type")
    .Output("output_handle: variant")
    .Attr("element_dtype: type")
    .Attr("shape_type: {int32, int64}")
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      DataType element_dtype;
      TF_RETURN_IF_ERROR(c->GetAttr("element_dtype", &element_dtype));
      shape_inference::ShapeHandle element_shape;
      TF_RETURN_IF_ERROR(c->MakeShapeFromShapeTensorTreatScalarAsUnknownShape(
          2, &element_shape));
      c->set_output_handle_shapes_and_types(0,
                                            {{element_shape, element_dtype}});
      c->set_output(0, c->Scalar());
      return Status::OK();
    });

REGISTER_OP("TensorListScatterV2")
    .Input("tensor: element_dtype")
    .Input("indices: int32")
    .Input("element_shape: shape_type")
    .Input("num_elements: int32")
    .Output("output_handle: variant")
    .Attr("element_dtype: type")
    .Attr("shape_type: {int32, int64}")
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      DataType element_dtype;
      TF_RETURN_IF_ERROR(c->GetAttr("element_dtype", &element_dtype));
      shape_inference::ShapeHandle element_shape;
      TF_RETURN_IF_ERROR(c->MakeShapeFromShapeTensorTreatScalarAsUnknownShape(
          2, &element_shape));
      c->set_output_handle_shapes_and_types(0,
                                            {{element_shape, element_dtype}});
      c->set_output(0, c->Scalar());
      return Status::OK();
    });

REGISTER_OP("TensorListScatterIntoExistingList")
    .Input("input_handle: variant")
    .Input("tensor: element_dtype")
    .Input("indices: int32")
    .Output("output_handle: variant")
    .Attr("element_dtype: type")
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      shape_inference::ShapeHandle ignored;
      // Check that tensor is at least a vector.
      TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(1), 1, &ignored));
      // Check that indices is a vector.
      TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 1, &ignored));

      DataType element_dtype;
      TF_RETURN_IF_ERROR(c->GetAttr("element_dtype", &element_dtype));
      shape_inference::ShapeHandle element_shape = c->UnknownShape();

      auto* handle_data = c->input_handle_shapes_and_types(0);
      if (handle_data != nullptr) {
        TF_RETURN_IF_ERROR(VerifyHandleData(c, *handle_data, element_dtype));
        element_shape = GetElementShapeFromHandleData(*handle_data);
      }
      c->set_output_handle_shapes_and_types(0,
                                            {{element_shape, element_dtype}});
      c->set_output(0, c->Scalar());
      return Status::OK();
    });

REGISTER_OP("TensorListConcatLists")
    .Input("input_a: variant")
    .Input("input_b: variant")
    .Attr("element_dtype: type")
    .Output("output: variant")
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      auto input_a = c->input(0);
      auto input_b = c->input(1);
      TF_RETURN_IF_ERROR(c->Merge(input_a, input_b, &input_a));
      c->set_output(0, input_a);

      DataType element_dtype;
      TF_RETURN_IF_ERROR(c->GetAttr("element_dtype", &element_dtype));

      auto* handle_data_a = c->input_handle_shapes_and_types(0);
      auto* handle_data_b = c->input_handle_shapes_and_types(1);
      bool handle_data_a_nonempty = handle_data_a && !handle_data_a->empty();
      bool handle_data_b_nonempty = handle_data_b && !handle_data_b->empty();
      if (!(handle_data_a_nonempty || handle_data_b_nonempty)) {
        c->set_output_handle_shapes_and_types(
            0, {{c->UnknownShape(), element_dtype}});
        return Status::OK();
      }
      shape_inference::ShapeAndType list_shape_type_a =
          handle_data_a_nonempty ? handle_data_a->at(0) : handle_data_b->at(0);
      const shape_inference::ShapeAndType& list_shape_type_b =
          handle_data_b_nonempty ? handle_data_b->at(0) : handle_data_a->at(0);
      if (list_shape_type_a.dtype != element_dtype) {
        return errors::InvalidArgument("input_a.type != element_dtype: ",
                                       DataTypeString(list_shape_type_a.dtype),
                                       " vs. ", DataTypeString(element_dtype));
      }
      if (list_shape_type_b.dtype != element_dtype) {
        return errors::InvalidArgument("input_b.type != element_dtype: ",
                                       DataTypeString(list_shape_type_b.dtype),
                                       " vs. ", DataTypeString(element_dtype));
      }
      TF_RETURN_IF_ERROR(c->Merge(list_shape_type_a.shape,
                                  list_shape_type_b.shape,
                                  &list_shape_type_a.shape));
      c->set_output_handle_shapes_and_types(0, {list_shape_type_a});
      return Status::OK();
    });

}  // namespace
}  // namespace tensorflow
