weight_row = array_ops.ones(
        shape=array_ops.shape(input_vocabs_placeholder), dtype=dtypes.float32
    )
    # shape: (?, 2)
    weight = array_ops.transpose_v2(
        array_ops_stack.stack([weight_row, weight_row])
    )
    # shape: (2, 2)
    output_tensor = math_ops.matmul(matmul_input, weight)

    return input_vocabs_placeholder, lookup_vals, output_tensor

// Returns a tensor of zeros with the same shape and type as x.
Status ZerosLike(AbstractContext* ctx, AbstractTensorHandle* const x,
                 AbstractTensorHandle** y, const char* name = nullptr,
                 const char* raw_device_name = nullptr);

// Returns the shape of a tensor.
Status Shape(AbstractContext* ctx, AbstractTensorHandle* const input,

}

// Converts the tensor shape proto into an MLIR shape attribute.
absl::StatusOr<mlir::Attribute> ConvertTensorShapeProto(
    const TensorShapeProto& shape, mlir::MLIRContext* context) {
  if (shape.unknown_rank())
    return mlir::TF::ShapeAttr::get(context, std::nullopt);

  llvm::SmallVector<int64_t, 4> dims;
  dims.reserve(shape.dim().size());
  for (const auto& dim : shape.dim()) {

// if (TF_GetCode(status) != TF_OK) {
//   // handle error
// }
//
// SHAPE INFERENCE
// ---------------
//
// You can provide a shape inference function that TensorFlow will call when it
// wants to understand the shape of outputs that the op will produce. Use the
// TF_OpDefinitionBuilderSetShapeInferenceFunction function to register a shape
// inference function pointer with TensorFlow. The following is an example of a

  attr {
    key: "_user_specified_name"
    value {
      s: "x"
    }
  }
  attr {
    key: "dtype"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "shape"
    value {
      shape {
        dim {
          size: 10
        }
      }
    }
  }
}
node {
  name: "y"
  op: "Placeholder"
  attr {
    key: "_user_specified_name"
    value {
      s: "y"

      return;
    }
    if (!TensorShape{e.tensor.shape()}.IsSameSize(
            TensorShape{first.tensor.shape()})) {
      status = errors::InvalidArgument(
          "Initializer list components should all have the same shape");
      return;
    }
  }

  // Form the new shape.
  TensorShape shape{static_cast<int64_t>(v.size())};
  shape.AppendShape(TensorShape{first.tensor.shape()});

  Tensor t(first.tensor.dtype(), shape);

#include "tsl/platform/statusor.h"

namespace tensorflow {

static size_t GetTensorSize(const TensorShape& shape, const DataType dtype) {
  return shape.num_elements() * DataTypeSize(dtype);
}

absl::StatusOr<Tensor> MakeTensorFromPjRtBuffer(
    const DataType dtype, const TensorShape& shape,
    std::unique_ptr<xla::PjRtBuffer> pjrt_buffer) {
  TF_ASSIGN_OR_RETURN(std::unique_ptr<xla::PjRtBuffer::ExternalReference> ref,

  func.return %0 : !tfr.tensor<K>
}

// -----

// CHECK-LABEL: get_shape
func.func @get_shape(%arg0: !tfr.tensor) -> (!shape.shape, !shape.shape) {
  %0 = tfr.get_shape %arg0 -> !shape.shape
  %1 = "tfr.get_shape"(%arg0) : (!tfr.tensor) -> !shape.shape
  func.return %0, %1 : !shape.shape, !shape.shape
}

// -----

// CHECK-LABEL: get_real_shape
func.func @get_real_shape(%arg0: tensor<1x2xf32>) -> tensor<2xindex> {

          cpu_tensor->shape(), cpu_tensor->dtype(), std::nullopt);
  TF_ASSIGN_OR_RETURN(xla::Shape shape,
                      shape_determination_fns.shape_representation_fn(
                          cpu_tensor->shape(), cpu_tensor->dtype(),
                          /*fast_mem=*/false, layout_preference));
  const xla::Layout* device_layout = &(shape.layout());
  // The device id should match the local_hardware_id in

  "func.func"() ({
  ^bb0(%arg0: tensor<i32>):
    "func.return"(%arg0) : (tensor<i32>) -> ()
  }) {sym_name = "main_while_body_4225150", sym_visibility = "private", tf._input_shapes = [#tf_type.shape<>], tf.signature.is_stateful, function_type = (tensor<i32>) -> (tensor<i32>)} : () -> ()