jax
a2880363 - Move NamedSharding.__eq__ and NamedSharding.__hash__ into C++.

Commit
316 days ago
Move NamedSharding.__eq__ and NamedSharding.__hash__ into C++. * change NamedSharding::spec_ to use a more precise type, now PartitionSpec is a type known to C++. * update PartitionSpec::Eq and PartitionSpec::Hash to be more usable from other C++ classes. * add a thread-safe cache for the hash value of a NamedSharding. * move the existing ShardingHash and ShardingEqual functions into jax_jit and rename them, since they are implementations that only really make sense for jit. PiperOrigin-RevId: 770635344
Author
Parents
Loading