// 'input_permutation' is a mapping from old argument numbers to new argument
// numbers, whereas 'output_permutation' is the same for outputs. Both
// 'input_permutation' and 'output_permutation' are initialized to the identity
// permutation. 'nodedef' is the NodeDef for the call to the function under
// construction, provided to allow additional attributes to be set.
// The rewrite may also change the NodeDef's operator name, and that

  //     ..., dn-1). For example, a 2-D matrix stored in row-major order would
  //     have traversal_order = (d0, d1).
  //   - if block sparse with a k-dimensional block (0 <= k <= n), the
  //     traversal_order has n + k elements. The first n elements are still a
  //     permutation of (d0, ..., dn-1). The lask k elements are a permutation

import org.gradle.test.preconditions.IntegTestPreconditions
import spock.lang.Issue

@Requires(value = IntegTestPreconditions.IsEmbeddedExecutor, reason = """
This test is very expensive due to the permutation testing.
Because it tests the internal state of the resolution engine, the Gradle execution model does not matter.
Se we run the tests only in embedded mode
""")

	// Shuffle really ought not be called with n that doesn't fit in 32 bits.
	// Not only will it take a very long time, but with 2³¹! possible permutations,
	// there's no way that any PRNG can have a big enough internal state to
	// generate even a minuscule percentage of the possible permutations.
	// Nevertheless, the right API signature accepts an int n, so handle it as best we can.
	for i := n - 1; i > 0; i-- {

  //     ..., dn-1). For example, a 2-D matrix stored in row-major order would
  //     have traversal_order = (d0, d1).
  //   - if block sparse with a k-dimensional block (0 <= k <= n), the
  //     traversal_order has n + k elements. The first n elements are still a
  //     permutation of (d0, ..., dn-1). The lask k elements are a permutation

template <class Op>
static LogicalResult InferenceFoldOperandsPermutation(
    ArrayRef<int64_t> permutation, Op* op) {
  // FusedBatchNorm in training mode is a layout sentitive operation, and should
  // have already assigned an optimal data format.
  if (op->getIsTraining()) return failure();
  return ::mlir::TF::FoldOperandsPermutation(permutation, op);
}

template <class Op>

  if (op_name == "stablehlo.slice" &&
      (field_name == "start_indices" || field_name == "limit_indices" ||
       field_name == "strides"))
    return true;
  if (op_name == "stablehlo.transpose" && field_name == "permutation")
    return true;
  return false;
}

class TflToStablehloPass
    : public mlir::PassWrapper<TflToStablehloPass,
                               mlir::OperationPass<mlir::func::FuncOp>> {
 public:

  //     ..., dn-1). For example, a 2-D matrix stored in row-major order would
  //     have traversal_order = (d0, d1).
  //   - if block sparse with a k-dimensional block (0 <= k <= n), the
  //     traversal_order has n + k elements. The first n elements are still a
  //     permutation of (d0, ..., dn-1). The lask k elements are a permutation

def : Pat<(TF_CastOp $arg, ConstBoolAttrFalse), (MHLO_ConvertOp $arg)>;

def : Pat<(TF_TransposeOp:$res $arg, (ConstantLikeMatcher ElementsAttr:$permutation)),
          (MHLO_TransposeOp $arg, (CastElementsToI64Elements $permutation))>;


// Lowering these ops with static shape to mhlo.reshape
foreach TfOp = [TF_ExpandDimsOp, TF_ReshapeOp, TF_SqueezeOp, ] in {

      List<T> expectedPermutation, Iterator<List<T>> permutations) {
    assertTrue("Expected another permutation, but there was none.", permutations.hasNext());
    assertEquals(expectedPermutation, permutations.next());
  }

  private <T> void assertNoMorePermutations(Iterator<List<T>> permutations) {
    assertFalse("Expected no more permutations, but there was one.", permutations.hasNext());
    try {
      permutations.next();