jax
654a2f6e - Strip all info except for shape and dtype while matching input and output avals for donation. For avals that are sharded, we match on shard_shape.

Commit

1 year ago

Strip all info except for shape and dtype while matching input and output avals for donation. For avals that are sharded, we match on shard_shape. PiperOrigin-RevId: 722917186

References

#26294 - Strip all info except for shape and dtype while matching input and output avals for donation. For avals that are sharded, we match on shard_shape.

#221 - CI: 02/04/25 upstream sync

Author

yashk2810

Committer

Google-ML-Automation

Parents

6281b860

jax 654a2f6e - Strip all info except for shape and dtype while matching input and output avals for donation. For avals that are sharded, we match on shard_shape.

jax
654a2f6e - Strip all info except for shape and dtype while matching input and output avals for donation. For avals that are sharded, we match on shard_shape.