--new_dim_count;
    }

    TF_Tensor* output = TF_AllocateTensor(k->output_data_type, dims, 0,
                                          TF_DataTypeSize(k->output_data_type));
    TF_TensorBitcastFrom(tensor, k->output_data_type, output, dims,
                         new_dim_count, status);
    if (TF_GetCode(status) == TF_OK) {
      TF_SetOutput(ctx, 0, output, status);
    }
    delete[] dims;

//   * TF_DataTypeSize(TF_TensorType(from)) != 0
//   * TF_DataTypeSize(type) != 0
//
// If any of the requirements are not met, *status is set to
// TF_INVALID_ARGUMENT.
TF_CAPI_EXPORT extern void TF_TensorBitcastFrom(const TF_Tensor* from,
                                                TF_DataType type, TF_Tensor* to,
                                                const int64_t* new_dims,

int64_t TF_TensorElementCount(const TF_Tensor* t) {
  int64_t result = 1;
  int rank = TF_NumDims(t);
  for (int dim = 0; dim < rank; ++dim) {
    result *= TF_Dim(t, dim);
  }
  return result;
}

void TF_TensorBitcastFrom(const TF_Tensor* from, TF_DataType type,
                          TF_Tensor* to, const int64_t* new_dims,
                          int num_new_dims, TF_Status* status) {
  TF_SetStatus(status, TF_OK, "");

  EXPECT_EQ(6 * TF_DataTypeSize(TF_UINT64), TF_TensorByteSize(a));
  EXPECT_EQ(TF_DataTypeSize(TF_UINT64), TF_TensorByteSize(b));

  int64_t new_dims[] = {3, 2};
  TF_Status* status = TF_NewStatus();
  TF_TensorBitcastFrom(a, TF_UINT64, b, new_dims, 2, status);
  ASSERT_EQ(TF_OK, TF_GetCode(status));
  TF_DeleteStatus(status);

  EXPECT_EQ(6, TF_TensorElementCount(a));
  EXPECT_EQ(6, TF_TensorElementCount(b));