Mirrors_Framework/deepmind-research
1
0
Fork 0
mirror of https://github.com/google-deepmind/deepmind-research.git synced 2026-06-02 14:45:25 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Replace jnp.DeviceArray with the new public type jax.Array.

Browse Source 
PiperOrigin-RevId: 477558995
This commit is contained in:
Yash Katariya
2022-09-28 21:53:07 +00:00
committed by Diego de las Casas
parent 58fb45db7e
commit da0f2de14d
1 changed files with 1 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tandem_dqn/agent.py
+1 -2
View File
@@ -338,8 +338,7 @@ class TandemDqn(parts.Agent):
def statistics(self) -> Mapping[Text, float]: def statistics(self) -> Mapping[Text, float]:
"""Returns current agent statistics as a dictionary.""" """Returns current agent statistics as a dictionary."""
# Check for DeviceArrays in values as this can be very slow. # Check for DeviceArrays in values as this can be very slow.
assert all( assert all(not isinstance(x, jax.Array) for x in self._statistics.values())
not isinstance(x, jnp.DeviceArray) for x in self._statistics.values())
return self._statistics return self._statistics
@property @property