ArrayRef<int64_t> shape = ranked_type.getShape();
    assert(permutation.size() == shape.size());

    SmallVector<int64_t, 4> new_shape(permutation.size());
    for (size_t i = 0; i < permutation.size(); ++i)
      new_shape[i] = shape[permutation[i]];

    return RankedTensorType::get(new_shape, ranked_type.getElementType());
  }

  return type;
}

bool AreCancellablePermutations(DenseIntElementsAttr perm0,

  XLA_Shape old_shape_c = {};
  XLA_Shape new_shape_c = {};
  TfTpu_ExecutorApiFn *executor = stream_executor::tpu::ExecutorApiFn();
  if (!stream_executor::tpu::IsInitialized(executor)) {
    return failure();
  }
  ApiConverter::ToC(old_shape, &old_shape_c);
  executor->TpuTransferManager_GetInfeedLayoutFn(&old_shape_c, &new_shape_c);
  xla::Shape new_shape = ApiConverter::FromC(&new_shape_c);
  ApiConverter::Destroy(&old_shape_c);

      "outputs": [1],
      "operators": [
        {
          "inputs": [ 0 ],
          "outputs": [ 1 ],
          "builtin_options_type": "ReshapeOptions",
          "builtin_options": {
            "new_shape": [ 2, 2 ]
          }
        }
      ]
    }
  ],
  "buffers": []

    %2 = "tf.PartitionedCall"(%1) {config = "", config_proto = "", executor_type = "", f = @get_shape} : (tensor<?x?xf32>) -> (tensor<?xi64>)

    // CHECK: return %[[RESULT]]
    func.return %2 : tensor<?xi64>
  }

  // CHECK-LABEL: func @get_shape
  func.func @get_shape(%arg0 : tensor<*xi64>) -> tensor<?xi64> {
    %0 = "tf.Shape"(%arg0) : (tensor<*xi64>) -> tensor<?xi64>

}

// cSHAKE specific context
type cshakeState struct {
	*state // SHA-3 state context and Read/Write operations

	// initBlock is the cSHAKE specific initialization set of bytes. It is initialized
	// by newCShake function and stores concatenation of N followed by S, encoded
	// by the method specified in 3.3 of [1].
	// It is stored here in order for Reset() to be able to put context into
	// initial state.
	initBlock []byte
}

    }

    // Save the new state value to the resource.
    pack_args.push_back(key);
    Value new_state = rewriter.create<PackOp>(loc, res_type, pack_args);
    rewriter.create<AssignVariableOp>(loc, rng_op.getResource(), new_state);

    // Pad the original state as necessary to fill the output shape.
    int pad = tensorflow::RNG_MAX_COUNTER_SIZE - counter_size;

    TF::ConstOp new_reshape_shape = GetI64ConstantTensor(
        rewriter, ArrayRef<int64_t>(new_reshape_dims), op.getLoc());
    auto new_reshape_type = RankedTensorType::get(new_reshape_dims, el_ty);
    ReshapeOp new_reshape =
        rewriter.create<ReshapeOp>(new_reshape_shape.getLoc(), new_reshape_type,
                                   op.getInput(), new_reshape_shape);
    TF::ConstOp new_broadcast_shape =

     kernel_initializer=tf.variance_scaling_initializer(),
     data_format=data_format)

    # Use the image size without space-to-depth transform as the input of conv0.
    batch_size, h, w, channel = inputs.get_shape().as_list()
    conv0.build([
     batch_size, h * space_to_depth_block_size, w * space_to_depth_block_size,
     channel // (space_to_depth_block_size**2)
    ])

    kernel = conv0.weights[0]