%split_dim_0 = arith.constant dense<0> : tensor<i32>
  %0, %1 = "tfl.split"(%split_dim_0, %arg0) {num_splits = 2 : i32} : (tensor<i32>, tensor<16x4xf32>) -> (tensor<8x4xf32>, tensor<8x4xf32>)
  %split_dim_1 = arith.constant dense<1> : tensor<i32>

  "operator_codes": [
    {
      "builtin_code": "SPLIT"
    }
  ],
  "subgraphs": [
    {
      "tensors": [
        {
          "shape": [
            1
          ],
          "name": "split_dim",
          "quantization": {
          }
        },
        {
          "shape": [
            256,
            32,
            32,
            3
          ],
          "name": "input",

  // TODO(renjieliu): change to use split_dim when we raise the constants
  // as well.
  int64_t split_dim = -1;
  for (int64_t d = 0; d < input_type.getRank(); ++d) {
    if (input_type.getDimSize(d) != output_type.getDimSize(d)) split_dim = d;
  }

  const SmallVector<Value, 4>& slice_outputs = SliceOutputs(
      split_op, input, input_type, split_dim, num_splits, &rewriter);

  Scope scope = Scope::NewRootScope();
  TensorShape x_shape({5, 2});
  auto x = Placeholder(scope, DT_DOUBLE, Placeholder::Shape(x_shape));
  // Split along the second dimension.
  auto split_dim = Const(scope, 1, {});
  auto y = Split(scope, split_dim, x, /* num_split */ 2);
  TensorShape y_shape = TensorShape({5, 1});
  double max_error;
  TF_ASSERT_OK((ComputeGradientError<double, double, double>(

  if (!input_type) return success();

  int64_t split_dim = split_dim_opt.value();
  const int64_t rank = input_type.getRank();
  if (split_dim < 0) split_dim += rank;
  if (split_dim < 0 || split_dim >= rank)
    return op.emitOpError("'split_dim' should be in [-rank, rank)");

  // If the 'split_dim' dimension of the 'input' tensor has a dynamic size,
  // there are no other checks.

      CompileTimeConstantOperand<TF::SliceOp, 1, 2>,       // $begin, $size
      CompileTimeConstantOperand<TF::SparseToDenseOp, 1>,  // $output_shape
      CompileTimeConstantOperand<TF::SplitOp, 0>,          // $split_dim
      // $size_splits, $split_dim
      CompileTimeConstantOperand<TF::SplitVOp, 1, 2>,
      CompileTimeConstantOperand<TF::StackV2Op, 0>,  // $max_size
      // $num_samples

// -----

func.func @testSplitNonConstSplitDim(%input: tensor<4x4xf32>, %split_dim: tensor<i32>) {
  %0:2 = "tf.Split"(%split_dim, %input) : (tensor<i32>, tensor<4x4xf32>) -> (tensor<*xf32>, tensor<*xf32>)
  func.return
}

func.func @testSplitUnknownRankSplitDim(%input: tensor<4x4xf32>, %split_dim: tensor<*xi32>) {
  %0:2 = "tf.Split"(%split_dim, %input) : (tensor<*xi32>, tensor<4x4xf32>) -> (tensor<*xf32>, tensor<*xf32>)

TEST_F(ArrayGradTest, SplitGrad) {
  TensorShape x_shape({5, 2});
  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
  // Split along the second dimension.
  auto split_dim = Const(scope_, 1, {});
  auto y = Split(scope_, split_dim, x, /* num_split */ 2);
  TensorShape y_shape = TensorShape({5, 1});
  RunTest({x}, {x_shape}, y.output, {y_shape, y_shape});
}

TEST_F(ArrayGradTest, SplitVGrad) {

  let description = [{
    Splits the `value` tensor along `split_dim` into a number of sub-tensors
    with same shape as the original one, except for `split_dim`. Same as
    tf.Split.
  }];

  let arguments = (ins
    TFL_TensorOf<[I32]>:$split_dim,
    TFL_TensorOf<[F32, I16, I32, I8, UI8, QI8, QUI8, QI16]>:$value,
    ConfinedAttr<I32Attr, [IntPositive]>:$num_splits

# tensorflow/core/api_def/base_api/*
# {op_name: {API's attribute name: OpDef's attribute name}}
_ATTRIBUTE_RENAMES = {
    'Mean': {'axis': 'reduction_indices'},
    'Split': {'axis': 'split_dim'},
    'SplitV': {'axis': 'split_dim'},
}


def _get_type_from_proto(arg_def=None, attr_def=None):
  if not arg_def:
    if attr_def.type == 'bool':
      return TFRTypes.I1
    elif attr_def.type == 'int32':