Add a benchmark for forward+backward for defuns.
Also fix some simple issues that I saw when I benchmarked it (goes from ~3500 examples/sec -> ~4000 examples/sec) - (nest) Expose is_mapping check that caches to python. - (nest) Stop calling flatten when unnecessary in pack_sequence_as - (nest) Set some functions to their swig wrappers directly (instead of wrapping them in another function) - Directly call the gen_math_ops call in _aggregate_grads to skip any unnecessary python overhead. - Stop falling back to slow path in _fast_fill. PiperOrigin-RevId: 209223633
Loading
Please sign in to comment