}

// 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;
}

            key: "_output_shapes"
            value {
              list {
              }
            }
          }
        }
        ret {
          key: "identity_1"
          value: "Identity_1:output:0"
        }
        attr {
          key: "_construction_context"
          value {
            s: "kEagerRuntime"
          }
        }
        attr {

    }
    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