Repository: tensorflow/tensorflow (language: cpp)

Issue:
Allow functions of `DistributedVariable`s to implement `custom_gradient`
Previously, variables created by functions decorated with `custom_gradient` had to pass:
```python 
isinstance(v, resource_variable_ops.ResourceVariable)
```
With this patch, they instead must pass:
```python
resource_variable_ops.is_resource_variable(v)
```
This allows functions that create [`DistributedVariable`](https://github.com/tensorflow/tensorflow/blob/a799b815e5b497fd85e36ec36e44df0f3cebf802/tensorflow/python/distribute/values.py#L429)s (which already implement the method [`_should_act_as_resource_variable`](https://github.com/tensorflow/tensorflow/blob/a799b815e5b497fd85e36ec36e44df0f3cebf802/tensorflow/python/distribute/values.py#L568)) to use `custom_gradient`, which is necessary for cross-replica batch normalization, among other things.


Discussion / hints:

Thanks for your pull request. It looks like this may be your first contribution to a Google open source project (if not, look below for help). Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

:memo: **Please visit <https://cla.developers.google.com/> to sign.**

Once you've signed (or fixed any issues), please reply here (e.g. `I signed it!`) and we'll verify it.

----

#### What to do if you already signed the CLA

##### Individual signers

*   It's possible we don't have your GitHub username or you're using a different email address on your commit. Check [your existing CLA data](https://cla.developers.google.com/clas) and verify that your [email is set on your git commits](https://help.github.com/articles/setting-your-email-in-git/).

##### Corporate signers

*   Your company has a Point of Contact who decides which employees are authorized to participate. Ask your POC to be added to the group of authorized contributors. If you don't know who your Point of Contact is, direct the Google project maintainer to [go/cla#troubleshoot](http://go/cla#troubleshoot) ([Public version](https://opensource.google.com/docs/cla/#troubleshoot)).
*   The email used to register you as an authorized contributor must be the email used for the Git commit. Check [your existing CLA data](https://cla.developers.google.com/clas) and verify that your [email is set on your git commits](https://help.github.com/articles/setting-your-email-in-git/).
*   The email used to register you as an authorized contributor must also be [attached to your GitHub account](https://github.com/settings/emails).
		

<!-- need_sender_cla -->
> Once you've signed (or fixed any issues), please reply here (e.g. `I signed it!`) and we'll verify it.

Signed the CLA
CLAs look good, thanks!

<!-- ok -->
@alextp, I'm not familiar with this code. Are you the right person to review this?
The change looks good, but I'd also love a unit test for it.
I would be happy to submit a unit test; however since there is not yet any `custom_gradient_test.py`, it seems befitting of a separate PR. 

Fix patch:
diff --git a/third_party/xla/xla/service/dump.cc b/third_party/xla/xla/service/dump.cc
index 3b238dfa4f4b77..7ffb1e18376ad8 100644
--- a/third_party/xla/xla/service/dump.cc
+++ b/third_party/xla/xla/service/dump.cc
@@ -886,20 +886,10 @@ std::string GetNonDefaultDebugOptions(const DebugOptions& debug_options) {
       continue;
     }
 
-    // Check if this field differs from default
-    if (reflection->HasField(debug_options, field) &&
-        !reflection->HasField(default_options, field)) {
-      // Field exists in debug_options but not defaults
+    if (GetValueAsString(reflection, debug_options, field) !=
+        GetValueAsString(reflection, default_options, field)) {
       absl::StrAppend(&non_default_options, field->name(), ": ",
                       GetValueAsString(reflection, debug_options, field), "\n");
-    } else if (reflection->HasField(debug_options, field)) {
-      // Field exists in both, compare values
-      if (GetValueAsString(reflection, debug_options, field) !=
-          GetValueAsString(reflection, default_options, field)) {
-        absl::StrAppend(&non_default_options, field->name(), ": ",
-                        GetValueAsString(reflection, debug_options, field),
-                        "\n");
-      }
     }
   }
 


Test patch:
diff --git a/third_party/xla/xla/service/dump_test.cc b/third_party/xla/xla/service/dump_test.cc
index 71599516d18a92..686c5ace5ee058 100644
--- a/third_party/xla/xla/service/dump_test.cc
+++ b/third_party/xla/xla/service/dump_test.cc
@@ -311,6 +311,7 @@ TEST(DumpTest, GetNonDefaultDebugOptions) {
   options.set_xla_gpu_enable_nccl_user_buffers(
       !default_options.xla_gpu_enable_nccl_user_buffers());
   options.set_xla_enable_dumping(true);
+  options.set_xla_gpu_enable_shared_constants(false);
   // Enum field
   options.clear_xla_gpu_enable_command_buffer();
   options.add_xla_gpu_enable_command_buffer(DebugOptions::CUBLAS);
@@ -350,6 +351,8 @@ TEST(DumpTest, GetNonDefaultDebugOptions) {
               100)));
   EXPECT_THAT(non_default_options,
               testing::HasSubstr("xla_gpu_enable_nccl_user_buffers: true"));
+  EXPECT_THAT(non_default_options,
+              testing::HasSubstr("xla_gpu_enable_shared_constants: false"));
   EXPECT_THAT(non_default_options,
               testing::HasSubstr("xla_gpu_enable_command_buffer: CUBLAS"));
   EXPECT_THAT(non_default_options,