def DecomposeResourceApplyMomentumOpNonNesterov :
  Pattern<
    (TF_ResourceApplyMomentumOp:$src_op
       $var_resource, $accum_resource, $lr, $grad, $momentum,
       BoolAttr:$_, ConstBoolAttrFalse:$use_nesterov
    ),
    [(TF_AddV2Op:$accum_new
      (TF_MulOp
        (CreateTFReadVariableOp $src_op, $grad, $accum_resource),
        $momentum
      ),
      $grad
     ),

  attr {
    key: "T"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "use_locking"
    value {
      b: true
    }
  }
  attr {
    key: "use_nesterov"
    value {
      b: false
    }
  }
}
node {
  name: "training/LossScaleOptimizer/SGD/update_159/ResourceApplyKerasMomentum"
  op: "ResourceApplyKerasMomentum"

  attr {
    key: "T"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "use_locking"
    value {
      b: true
    }
  }
  attr {
    key: "use_nesterov"
    value {
      b: false
    }
  }
}
node {
  name: "training/SGD/SGD/update_fc1000/kernel/ResourceApplyKerasMomentum"
  op: "ResourceApplyKerasMomentum"

    // CHECK: [[VAR_NEW:%.*]] = "tf.Sub"([[VAR]], [[ACCUM_NEW_LR]])
    // CHECK: "tf.AssignVariableOp"([[VAR_HANDLE]], [[VAR_NEW]])
    "tf.ResourceApplyMomentum"(%0, %1, %arg0, %arg1, %arg2) {use_locking = false, use_nesterov = false} : (tensor<!tf_type.resource<tensor<f32>>>, tensor<!tf_type.resource<tensor<f32>>>, tensor<f32>, tensor<f32>, tensor<f32>) -> ()
    tf_device.return
  }) : () -> ()
  func.return
}

    TF_Float32Tensor:$epsilon,
    TF_Float32Tensor:$accumulator,
    TF_Float32Tensor:$momentum,
    TF_Float32Tensor:$embedding_table,

    I64Attr:$feature_width,
    BoolAttr:$use_nesterov,
    F32Attr:$beta_2,
    F32Attr:$exponent
  );

  let results = (outs
    TF_Float32Tensor:$updated_embedding_table,
    TF_Float32Tensor:$updated_accumulator,

    DefaultValuedOptionalAttr<BoolAttr, "false">:$use_nesterov
  );

  let results = (outs);

  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<2>;
}

def TF_ResourceApplyMomentumOp : TF_Op<"ResourceApplyMomentum", []> {
  let summary = "Update '*var' according to the momentum scheme.";

  let description = [{
Set use_nesterov = True if you want to use Nesterov momentum.