Add support for custom mapping types to util.nest

Analagous to the existing support for custom collections.Sequence types. They need to be constructable with the same arguments as the base type for pack_sequence_as to work.

Leaves PyDict_* calls for dict subclasses, but adds more general (and likely much slower) fallbacks for instances of collections.Mapping which are not dict subclasses.

My hope is that this support will be enough so I can use a wrapper around dicts which does not inherit from dict in __setattr__ tracking (some tests failed without it). Inheriting from dict and properly shadowing a real dict seems impossible with CPython (since to shadow without synchronization issues, the wrapper needs to respond to updates to the original dict, but to work with e.g. {}.update(dict_subclass) the wrapper's C storage needs to also be updated).

PiperOrigin-RevId: 205858082

1 files changed, 1 insertions, 0 deletions

diff --git a/tensorflow/python/BUILD b/tensorflow/python/BUILD
index b5a0051c28..a7c60f5450 100644
--- a/tensorflow/python/BUILD
+++ b/tensorflow/python/BUILD
@@ -3408,6 +3408,7 @@ py_test(
         ":math_ops",
         ":util",
         "//third_party/py/numpy",
+        "@absl_py//absl/testing:parameterized",
     ],
 )