padding_value, CreateScalarConstValue<int8_t>(builder, loc, 0));

    SmallVector<int64_t> shape_add(rank, 0);
    shape_add[pack_dim] = 1;
    shape_value = builder.create<TF::AddOp>(
        loc, shape_type, shape_value,
        CreateConstValue<int64_t>(builder, loc, {rank}, shape_add));
  }
  packed_shape[pack_dim] /= 2;
  SmallVector<int64_t> divisor(rank, 1);
  divisor[pack_dim] = 2;

  std::vector<int64_t>* shape_arr = &tf_dlm_tensor_ctx->shape;
  std::vector<int64_t>* stride_arr = &tf_dlm_tensor_ctx->strides;
  shape_arr->resize(ndim);
  stride_arr->resize(ndim, 1);
  for (int i = 0; i < ndim; i++) {
    (*shape_arr)[i] = tensor->dim_size(i);
  }
  for (int i = ndim - 2; i >= 0; --i) {
    (*stride_arr)[i] = (*shape_arr)[i + 1] * (*stride_arr)[i + 1];
  }

  tensorflow::Scope s = tensorflow::Scope::NewRootScope();
  Output shape_a = ops::Placeholder(s.WithOpName("shape_a"), DT_INT32,
                                    ops::Placeholder::Attrs{});
  Output shape_b = ops::Placeholder(s.WithOpName("shape_b"), DT_INT32,
                                    ops::Placeholder::Attrs{});
  Output shape = ops::Add(s.WithOpName("shape"), shape_a, shape_b);

  const TF_TensorSpec* tensor_spec_a = TF_SignatureDefParamTensorSpec(param_a);
  const TF_Shape* shape_a = TF_TensorSpecShape(tensor_spec_a);

  // Input "a" is a scalar, float32 tensor
  EXPECT_EQ("a", std::string(TF_SignatureDefParamName(param_a)));
  EXPECT_EQ(TF_FLOAT, TF_TensorSpecDataType(tensor_spec_a));
  EXPECT_EQ(0, TF_ShapeDims(shape_a));

  const TF_SignatureDefParam* param_b = TF_SignatureDefParamListGet(args, 1);

// Shape op legalization.
//===----------------------------------------------------------------------===//

// -----

// CHECK-LABEL: func @shape_1D
func.func @shape_1D(%arg0: tensor<?xf32>) -> tensor<1xi32> {
  // CHECK: [[SHAPE:%.+]] = shape.shape_of %arg0
  // CHECK: [[TENSOR:%.+]] = arith.index_cast [[SHAPE]] : tensor<1xindex> to tensor<1xi32>