// CHECK-LABEL: avg_pool2d_3
// CHECK: %cst = arith.constant dense<[0, 2, 3, 1]> : tensor<4xi32>
// CHECK: %0 = "tfl.transpose"(%arg0, %cst) : (tensor<1x1x1x8xf32>, tensor<4xi32>) -> tensor<1x1x8x1xf32>
// CHECK{LITERAL}: %cst_0 = arith.constant dense<[[0, 0], [0, 0], [1, 1], [0, 0]]> : tensor<4x2xi32>

  func.return %1 : tensor<2x1xf32>

// CHECK:  %cst = arith.constant
// CHECK:  %0 = "tf.Reshape"(%arg0, %cst)
// CHECK-SAME: tensor<2x1xf32>
// CHECK:  %1 = "tf.FakeQuantWithMinMaxVars"(%0, %arg1, %arg2)
// CHECK:  return %1
}

// CHECK-LABEL: fakeQuantFollowedByReshapes

  func.return %1 : tensor<2x1xf32>

// CHECK:  %cst = arith.constant
// CHECK:  %0 = "tf.Reshape"(%arg0, %cst)
// CHECK-SAME: tensor<2x1xf32>
// CHECK:  %1 = "tf.FakeQuantWithMinMaxVars"(%0, %arg1, %arg2)
// CHECK:  return %1
}

// CHECK-LABEL: fakeQuantFollowedByReshapes

  // CHECK: %1 = "tfl.transpose"(%0, %[[cst_0]]) : (tensor<7x6x5x4x3x2xf32>, tensor<6xi32>) -> tensor<2x3x4x5x6x7xf32>
  // CHECK: %2 = "tfl.reshape"(%1, %[[cst_1]]) : (tensor<2x3x4x5x6x7xf32>, tensor<7xi32>) -> tensor<1x2x3x4x5x6x7xf32>
  // CHECK: return %2
}

// CHECK-LABEL: OptimizeTranposeWithRank7orMoreEffectiveRank4
func.func @OptimizeTranposeWithRank7orMoreEffectiveRank4(%arg0: tensor<56x8x56x1x1x1x7xf32> ) -> (tensor<1x1x8x56x56x7x1xf32>)  {

  func.return %4#0, %4#1 : tensor<i16>, tensor<i32>

  // CHECK:  %0 = mhlo.constant dense<false> : tensor<i1>
  // CHECK:  %1 = "mhlo.iota"() <{iota_dimension = 0 : i64}> : () -> tensor<2xi32>
  // CHECK-DAG:  %2 = mhlo.constant dense<32767> : tensor<i16>
  // CHECK:  %3 = mhlo.constant dense<0> : tensor<i32>
  // CHECK:  %cst = arith.constant dense<0> : tensor<1xi32>

// CHECK:         %0 = "tf.Sub"(%arg0, %arg1) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
// CHECK:         %1 = "tf.Mul"(%0, %0) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
// CHECK:         return %1 : tensor<4xf32>
// CHECK:       }

// -----

func.func @unrollSplit(%arg0: tensor<i32>, %arg1: tensor<1x8x8x1024xf32>) -> (tensor<1x8x8x256xf32>, tensor<1x8x8x256xf32>, tensor<1x8x8x256xf32>) {

def ConvertTFEinsumToXLADotV2Op : Pat<
  (TF_EinsumOp:$einsum
    $args, $equation),
  (CreateXlaDotV2OpFromTfEinsumOp
    $equation, $args, $einsum),
  [(IsInt32ElementType $einsum),
   // Constraint to check:
   // 1. The einsum has two inputs and one output.
   // 2. The einsum is not created by the convert function itself.
   // 3. Both inputs are int32 tensor.

	SetDefault(logger)

	// Calls to check must be one line apart.
	// Determine line where calls start.
	f, _ := runtime.CallersFrames([]uintptr{callerPC(2)}).Next()
	startLine = f.Line + 4
	// Do not change the number of lines between here and the call to check(0).

	logger.Log(ctx, LevelInfo, "")
	check(0)
	logger.LogAttrs(ctx, LevelInfo, "")
	check(1)
	logger.Debug("")
	check(2)
	logger.Info("")
	check(3)

	// start with a clean slate (check.Files may be called multiple times)
	check.files = nil
	check.imports = nil
	check.dotImportMap = nil

	check.firstErr = nil
	check.methods = nil
	check.untyped = nil
	check.delayed = nil
	check.objPath = nil
	check.cleaners = nil

	// determine package name and collect valid files
	pkg := check.pkg
	for _, file := range files {

// via action.describef for debugging, if desired.
func (check *Checker) later(f func()) *action {
	i := len(check.delayed)
	check.delayed = append(check.delayed, action{f: f})
	return &check.delayed[i]
}

// push pushes obj onto the object path and returns its index in the path.
func (check *Checker) push(obj Object) int {
	check.objPath = append(check.objPath, obj)
	return len(check.objPath) - 1
}