jax
8d304fe7 - Simplify the jaxpr emitted for jnp.repeat in the case that the input dim size is 1.

Commit

1 year ago

Simplify the jaxpr emitted for jnp.repeat in the case that the input dim size is 1. In this case, the repeat() operation is just a broadcast. PiperOrigin-RevId: 780242309

References

#30024 - Simplify the jaxpr emitted for jnp.repeat in the case that the input dim size is 1.

Author

hawkinsp

Committer

Google-ML-Automation

Parents

3c6903d6

jax 8d304fe7 - Simplify the jaxpr emitted for jnp.repeat in the case that the input dim size is 1.

jax
8d304fe7 - Simplify the jaxpr emitted for jnp.repeat in the case that the input dim size is 1.