// CHECK-DAG:       [[VAL_11:%.*]] = "tf.Const"() <{value = dense<0> : tensor<i32>}> : () -> tensor<i32>
// CHECK:           [[VAL_12:%.*]]:4 = "tf.SplitV"([[VAL_7]], [[VAL_10]], [[VAL_11]]) : (tensor<40x8xf32>, tensor<4xi32>, tensor<i32>) -> (tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>)

  // CHECK: "tfl.space_to_batch_nd"
}

func.func @split(%arg0: tensor<i32>, %arg1: tensor<1x4x3x3xf32>) -> tensor<1x4x3xf32> {
  %0:3 = "tf.Split"(%arg0, %arg1) : (tensor<i32>, tensor<1x4x3x3xf32>) -> (tensor<1x4x3xf32>, tensor<1x4x3xf32>, tensor<1x4x3xf32>)
  func.return %0#0 : tensor<1x4x3xf32>

  // CHECK-LABEL: split

//===----------------------------------------------------------------------===//
// SplitOp
//===----------------------------------------------------------------------===//

// Verifies the input and split dimension operands for tf.Split/tf.SplitV.
// Writes the split dimension's index (adjusted with input rank) via `dim_index`
// if it's a constant.
template <class Op>
LogicalResult VerifySplitInputAndSplitDim(Op op,

    SameOperandsAndResultsScale]> {
  let summary = "Splits a tensor into `num_split` tensors along one dimension.";

  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`. The grouping
    of the resultant sub-tensors is decided by `size-splits`. Same as tf.SplitV.
  }];

  let arguments = (ins

		},
		workloadAgnostic: true,
	})

	splits := [][]int{
		{50, 25, 25},
		{80, 10, 10},
	}
	if skipVM {
		splits = [][]int{
			{50, 50},
			{80, 20},
		}
	}
	for _, split := range splits {
		split := split
		t.RunTraffic(TrafficTestCase{
			name:           fmt.Sprintf("shifting-%d", split[0]),

      // CHECK:      %[[SPLIT_0_OUT:[a-z0-9]+]]:2 = "tf.Split"(%[[CONST_SPLIT_0_DIM]], %[[RI_0]])
      // CHECK:      %[[CONST_SPLIT_1_DIM:.*]] = "tf.Const"()
      // CHECK:      %[[SPLIT_1_OUT:[a-z0-9]+]]:2 = "tf.Split"(%[[CONST_SPLIT_1_DIM]], %[[SPLIT_0_OUT]]#0)
      // CHECK:      %[[CONST_SPLIT_2_DIM:.*]] = "tf.Const"()
      // CHECK:      %[[SPLIT_2_OUT:[a-z0-9]+]]:2 = "tf.Split"(%[[CONST_SPLIT_2_DIM]], %[[SPLIT_0_OUT]]#1)

  %4, %5 = "tfl.split"(%split_dim_2, %arg0) {num_splits = 2 : i32} : (tensor<1xi32>, tensor<16x4xf32>) -> (tensor<16x2xf32>, tensor<16x2xf32>)
  %6:2 = "tfl.split"(%split_dim_2, %arg0) {num_splits = 2 : i32} : (tensor<1xi32>, tensor<16x4xf32>) -> (tensor<16x2xf32>, tensor<16x?xf32>)
  %7:2 = "tfl.split"(%split_dim_2, %arg0) {num_splits = 2 : i32} : (tensor<1xi32>, tensor<16x4xf32>) -> (tensor<?x2xf32>, tensor<16x?xf32>)

def TPUColocateSplitsPass : Pass<"tf-tpu-colocate-splits", "mlir::func::FuncOp"> {
  let summary = "Colocates each Split op with its predecessor";
  let constructor = "TFTPU::CreateTPUColocateSplitsPass()";
  let description = [{
    It is beneficial for performance to assign a `Split` op to the same device
    as its predecessor. This is because the weight of cut edges is always
    minimized when the `Split` is with its predecessor. This colocation

//
// The environment variable must be quoted correctly for
// quoted.Split. This should be done before building
// anything, for example, in BuildInit.
func envList(key, def string) []string {
	v := cfg.Getenv(key)
	if v == "" {
		v = def
	}
	args, err := quoted.Split(v)
	if err != nil {
		panic(fmt.Sprintf("could not parse environment variable %s with value %q: %v", key, v, err))

func validateGatewayNames(gatewayNames []string, gatewaySemantics bool) (errs Validation) {
	for _, gatewayName := range gatewayNames {
		parts := strings.SplitN(gatewayName, "/", 2)
		if len(parts) != 2 {
			if strings.Contains(gatewayName, ".") {
				// Legacy FQDN style
				parts := strings.Split(gatewayName, ".")
				recommended := fmt.Sprintf("%s/%s", parts[1], parts[0])
				errs = AppendValidation(errs, WrapWarning(fmt.Errorf(