//
//   Slice(op, begin, size <must be constant>) =>
//     Slice(op, begin, actual_size(op.shape(), size, begin));
//       _XlaCompileTimeConstantInputs={2}
//
// where
//
//   actual_size(op_shape, size, begin)[i] =
//     size[i] == -1 ? (op_shape[i] - size[i])
//                   : size[i]
//
// This pass, combined with jit/partially_decluster_pass, reduces the number of

  std::vector<int32_t> data{1, 2, 3, 4, 5, 6};
  xla::Shape xla_shape = xla::ShapeUtil::MakeShape(xla::S32, dimensions);
  TF_ASSERT_OK_AND_ASSIGN(
      auto pjrt_buffer,
      pjrt_client->BufferFromHostBuffer(
          data.data(), xla_shape.element_type(), xla_shape.dimensions(),
          /*byte_strides=*/std::nullopt,

      has_bias: bool,
      use_kernel: bool,
  ):
    n = 5
    x_shape = [v if v is not None else n for v in shapes[0]]
    y_shape = [v if v is not None else n for v in shapes[1]]

    class MatmulModel(module.Module):

      def __init__(self, bias: Optional[core.Tensor]):
        self._bias = bias
        self._kernel = np.random.uniform(size=y_shape).astype('f4')
        self._min = (-0.8, -0.8, -0.9)

#define TENSORFLOW_C_TF_SHAPE_INTERNAL_H_

#include "tensorflow/c/conversion_macros.h"
#include "tensorflow/core/framework/tensor_shape.h"

typedef struct TF_Shape TF_Shape;

namespace tensorflow {

DEFINE_CONVERSION_FUNCTIONS(tensorflow::PartialTensorShape, TF_Shape);

}

  std::vector<int64_t> in_shape{input_type.getShape().vec()};
  std::vector<int64_t> out_shape{output_type.getShape().vec()};

  // If the reshape changes the number of dimensions so it cannot be interpreted
  // as a transpose.
  if (in_shape.size() != out_shape.size()) {
    return false;
  }

  in_shape.erase(std::remove(in_shape.begin(), in_shape.end(), 1),
                 in_shape.end());

#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_PUBLIC_TENSOR_SPEC_H_

#include <stddef.h>

#include "tensorflow/c/c_api_macros.h"
#include "tensorflow/c/tf_datatype.h"
#include "tensorflow/c/tf_shape.h"

#ifdef __cplusplus
extern "C" {
#endif  // __cplusplus

// An opaque type corresponding to TensorSpec
typedef struct TF_TensorSpec TF_TensorSpec;

// Returns the dtype associated with the TensorSpec.

  mlir::FunctionType func_type = main_fn.getFunctionType();
  for (auto input_type : func_type.getInputs()) {
    tensorflow::TensorShape tensor_shape;
    xla::Shape xla_shape = xla::TypeToShape(input_type);
    TF_RETURN_IF_ERROR(tensorflow::TensorShape::BuildTensorShape(
        xla_shape.dimensions(), &tensor_shape));
    arg_shapes.emplace_back(tensor_shape);

    DataType dtype;
    TF_RETURN_IF_ERROR(ConvertToDataType(input_type, &dtype));

#include "tensorflow/c/tf_shape_internal.h"

extern "C" {

TF_DataType TF_TensorSpecDataType(const TF_TensorSpec* spec) {
  return static_cast<TF_DataType>(tensorflow::unwrap(spec)->dtype());
}

const TF_Shape* TF_TensorSpecShape(const TF_TensorSpec* spec) {
  return tensorflow::wrap(&tensorflow::unwrap(spec)->shape());
}

// Represents a (partially-defined) shape.
typedef struct TF_Shape {
  int num_dims;  // Must be >= -1; -1 represents unknown rank.
  int64_t* dim_sizes;
} TF_Shape;

// Add a new parameter to a TensorFlow Function.
TF_AbstractTensor* TF_AddFunctionParameter(TF_ExecutionContext* func,
                                           TF_DataType dtype, TF_Shape shape,
                                           TF_Status* s);

      [shape_determination_fns](
          const xla::Shape& xla_shape) -> absl::StatusOr<xla::Shape> {
    TensorShape shape;
    TF_RETURN_IF_ERROR(XLAShapeToTensorShape(xla_shape, &shape));
    TF_ASSIGN_OR_RETURN(DataType dtype, EncodePrimitiveTypeAsDataType(
                                            xla_shape.element_type()));
    auto layout_preference = shape_determination_fns.layout_preference_fn(