Make sure strategy.run.reduce(SUM) works across all axis.
Make sure strategy.run.reduce(SUM) works across all axis.

PiperOrigin-RevId: 375820623
Change-Id: I17c88767dd30d80faed9e7bdd531a5e74e5a0bde

tensorflow/agents,tensorflow/agents