search

google-deepmind / acme

A library of reinforcement learning components and agents
Apache License 2.0
3.52k stars 426 forks source link

Support JAX Array #281

Open cemlyn007 opened 1 year ago

cemlyn007 commented 1 year ago

Happy New Year guys!

Are you guys planning on supporting JAX Array and removing all uses of DeviceArray, ShardedDeviceArray, and GlobalDeviceArray types in JAX?