jax
408bd93e - Allow passing an aval directly to TypedNdArray.

Commit

10 days ago

Allow passing an aval directly to TypedNdArray. This saves reconstructing it if we already have it. Don't redundantly store weak_type, get it from the aval if needed. PiperOrigin-RevId: 899295373

References

#36746 - Allow passing an aval directly to TypedNdArray.

#36909 - [ROCm] Skip unit tests in `testEighIdentity` because of hipSolver issue.

#37071 - [ROCm] Switch to hermetic llvm

#37072 - [ROCm] Update rules_ml_toolchain reference

#37085 - Upgrade upstream ROCm CI from 7.2.0 to 7.2.2

Author

hawkinsp

Committer

Google-ML-Automation

Parents

8f39e8fc

jax 408bd93e - Allow passing an aval directly to TypedNdArray.

jax
408bd93e - Allow passing an aval directly to TypedNdArray.