jax

837f78e5 - Set cotangent sharding to full replicated if the corresponding primal sharding is empty in dot transpose.