flax
Use `jnp.stack` instead of `np.stack` in `flax.training.common_utils.stack_forest`
#4991
Merged
Merged
Use `jnp.stack` instead of `np.stack` in `flax.training.common_utils.stack_forest` #4991
copybara-service
merged
1 commit into
main
from
use-jnp-stack-in-stack_forest
use jnp.stack instead of np.stack
c08e60b3
samanklesaria
approved these changes
on 2025-09-30
cgarciae
approved these changes
on 2025-10-01
Login to write a write a comment.
Login via GitHub