search

google-deepmind / rlax

https://rlax.readthedocs.io
Apache License 2.0
1.23k stars 85 forks source link

Expose `use_jnp_split` arg in `tree_split_leaves`. #67

Closed copybara-service[bot] closed 2 years ago

copybara-service[bot] commented 2 years ago

Expose use_jnp_split arg in tree_split_leaves.

jnp.split shoud be preferred over np.split in jitted code as it can significantly reduce memory footprint.