jax

098624dc - Make the type checker match the runtime behavior of PartitionSpec not inherting from a tuple.