# PRESERVE:           %[[IDENTITY_N:.*]]:3, %{{.*}} = tf_executor.island(%[[ASSIGN_ADD_CTRL]]) wraps "tf.IdentityN"
# PRESERVE:           %[[ADD:.*]], %{{.*}} = tf_executor.island wraps "tf.Add"(%[[ARG_0]], %[[ARG_1]])

  } else if (yield.getNumOperands() == 1) {
    Value operand = yield.getOperand(0);
    auto identity = builder.create<TF::IdentityOp>(island.getLoc(),
                                                   operand.getType(), operand);
    yield.setOperand(0, identity.getOutput());
  } else {
    auto identity_n = builder.create<TF::IdentityNOp>(
        island.getLoc(), yield.getOperandTypes(), yield.getOperands());

namespace mlir {
namespace TFTPU {

namespace {

// This pass removes Identity/IdentityN ops from the TPU computation and
// reachable functions.
// TODO(lyandy): Remove this pass once resource op lifting is migrated to use
// resource alias analysis and support region based control flow. Removing
// Identity ops may remove `_XlaSharding` annotation attribute if Identity ops
// are used to propagate such information.

//
//   ```python
//   with tf.get_default_graph().gradient_override_map(
//       {'IdentityN': 'OverrideGradientWithG'}):
//     y, _ = identity_n([f(x), x])
//
//   @tf.RegisterGradient('OverrideGradientWithG')
//   def ApplyG(op, dy, _):
//     return [None, g(dy)]  # Do not backprop to f(x).
//   ```
Status IdentityN(AbstractContext* ctx,
                 absl::Span<AbstractTensorHandle* const> input,

	}

	// Initialize
	li, err := cr.NewCustomTokenCredentials(stsEndpoint, token, roleArn, opts...)
	if err != nil {
		log.Fatalf("Error initializing CustomToken Identity: %v", err)
	}

	v, err := li.Get()
	if err != nil {
		log.Fatalf("Error retrieving STS credentials: %v", err)
	}

	if displayCreds {
		fmt.Println("Only displaying credentials:")

    %0:3 = tf_executor.island {
      %1:2 = "tf.IdentityN"(%arg0, %arg1) : (tensor<i32>, tensor<f32>) -> (tensor<i32>, tensor<f32>) loc("identity@some_function")
      %2:2 = "tf.IdentityN"(%1#0, %1#1) : (tensor<i32>, tensor<f32>) -> (tensor<i32>, tensor<f32>) loc("identity@some_function")
      tf_executor.yield %2#0, %2#1 : tensor<i32>, tensor<f32> loc("identity@some_function")
    } loc("identity@some_function")
    tf_executor.fetch
  }

  %3 = "tf.IdentityN"(%2) {device = ""} : (tensor<1x2xf32>) -> tensor<1x2xf32>
  return %3 : tensor<1x2xf32>
}

// CHECK: func @cast_bf16_matmul_to_fp32
// CHECK-DAG: %[[cst:.*]] = "tf.Const"() <{value = dense<{{.*}}> : tensor<10x2xf32>}> : () -> tensor<10x2xf32>
// CHECK: %[[matmul:.*]] = "tf.MatMul"(%arg0, %[[cst]])
// CHECK: %[[identity:.*]] = "tf.IdentityN"(%[[matmul]])

    }
  }
}
node {
  name: "Identity"
  op: "Identity"
  input: "cond"
  device: "/device:TPU_REPLICATED_CORE:0"
  attr {
    key: "T"
    value {
      type: DT_INT32
    }
  }
  attr {
    key: "_tpu_replicate"
    value {
      s: "cluster"
    }
  }
}
node {
  name: "Identity_1"
  op: "Identity"
  input: "cond:1"

// RUN: tf-opt %s -tf-tpu-identity-pruning | FileCheck %s --dump-input=always

// Tests Identity op in cluster is pruned away.

// CHECK-LABEL: func @testIdentity
// CHECK-SAME: ([[ARG0:%.*]]: tensor<i32>)
func.func @testIdentity(%arg0: tensor<i32>) {
  // CHECK-NOT:  "tf.Identity"
  // CHECK:      "tf_device.cluster"
  // CHECK-NEXT: tf_device.return [[ARG0]]
  %0 = "tf_device.cluster"() ({
    %1 = "tf.Identity"(%arg0) : (tensor<i32>) -> tensor<i32>

  // CHECK-NEXT: bb0(%[[ARG_3:.*]]: tensor<i32>, %[[ARG_4:.*]]: tensor<*xf32>)
  ^bb0(%arg3: tensor<i32>, %arg4: tensor<*xf32>):
    // CHECK-NEXT: %[[RESULT_1:.*]] = "tf.Identity"(%[[ARG_3]])
    %1 = "tf.Identity"(%arg3) : (tensor<i32>) -> tensor<i32>
    %2 = "tf.Identity"(%arg1) : (tensor<i32>) -> tensor<i32>
    // CHECK-NEXT: %[[RESULT_2:.*]] = "tf.NotEqual"(%[[RESULT_1]], %[[ARG_1]])