refactor: use jnp.take_along_axis for return_last_only indexing#588
Open
stanley1208 wants to merge 1 commit intogoogle-deepmind:mainfrom
Open
refactor: use jnp.take_along_axis for return_last_only indexing#588stanley1208 wants to merge 1 commit intogoogle-deepmind:mainfrom
stanley1208 wants to merge 1 commit intogoogle-deepmind:mainfrom
