Refactor drjax primitives to use jnp.broadcast_to and improve explicit sharding support.#37

copybara-service[bot] merged 1 commit into

maingoogle/drjax:mainfrom

cl/820412444google/drjax:cl/820412444Copy head branch name to clipboard