fix(batch_norm): allow string cross_replica_axis for vmap support#859
Open
aditya161205 wants to merge 2 commits intogoogle-deepmind:mainfrom
Open
fix(batch_norm): allow string cross_replica_axis for vmap support#859aditya161205 wants to merge 2 commits intogoogle-deepmind:mainfrom
aditya161205 wants to merge 2 commits intogoogle-deepmind:mainfrom
