Function _unpmap

diffmimic/brax_lib/agent_diffmimic.py:57–58 · view source on GitHub ↗

(v)

Source from the content-addressed store, hash-verified

55
56
57	def _unpmap(v):
58	return jax.tree_util.tree_map(lambda x: x[0], v)
59
60
61	def _data_pmap(v, local_devices_to_use):

trainFunction · 0.85

no outgoing calls

no test coverage detected