tf_executor.yield %identity : tensor<i32>
        }
        %island_1_result, %island_1_control = tf_executor.island {
          %identity_n:2 = "tf.IdentityN"(%arg1, %island_0_result) : (tensor<i32>, tensor<i32>) -> (tensor<i32>, tensor<i32>)
          tf_executor.yield %identity_n#0
        }
        tf_executor.fetch %island_0_result, %island_1_result : tensor<i32>, tensor<i32>
      }

}

// Trace the graph to find out the actual operation.
Value getActualValue(Operation *op) {
  if (isa<TF::CastOp>(op)) {
    op = op->getOperand(0).getDefiningOp();
  }

  if (isa<TF::IdentityOp>(op)) {
    op = op->getOperand(0).getDefiningOp();
  }
  return op->getResult(0);
}

Value CreateXlaDotV2OpFromTfEinsumOp(OpBuilder &builder, Location loc,

}

// CHECK-LABEL: @FuseReshapeAroundBMMRHS
func.func @FuseReshapeAroundBMMRHS(%arg0: tensor<1x3x6x5x1024xf32>) -> tensor<1x3x6x5x8192xf32> attributes {tf.entry_function = {control_outputs = "", inputs = "inputs", outputs = "Identity_1"}} {
  %cst = arith.constant dense_resource<__elided__> : tensor<1x1024x8192xf32>
  %cst_0 = arith.constant dense_resource<__elided__> : tensor<5xi32>
  %cst_1 = arith.constant dense_resource<__elided__> : tensor<3xi32>

namespace mlir {
namespace TF {

namespace {
// Returns the equivalent Value skipping through identity nodes.
Value LookThroughIdentity(Value result) {
  while (isa_and_nonnull<IdentityOp, IdentityNOp>(result.getDefiningOp())) {
    auto op_result = result.cast<OpResult>();
    result = op_result.getOwner()->getOperand(op_result.getResultNumber());
  }
  return result;
}

    }
    return success();
  }
};

// Bypasses IdentityN op.
class ConvertIdentityNOp : public OpRewritePattern<TF::IdentityNOp> {
 public:
  using OpRewritePattern<TF::IdentityNOp>::OpRewritePattern;
  LogicalResult matchAndRewrite(TF::IdentityNOp op,
                                PatternRewriter &rewriter) const override {
    rewriter.replaceOp(op, op.getOperands());
    return success();

such that dx = g(dy). In Python,

```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).
```
  }];

  let arguments = (ins