{"record":{"id":"9a59b0bbdd7d5d08","repo":"jax-ml/jax","slug":"expected-unreduced-kind-to-be-of-type-jax-shardin-9a59b0","errorCode":null,"errorMessage":"Expected unreduced_kind to be of type `jax.sharding.UnreducedKind` but got {type(unreduced_kind)}","messagePattern":"Expected unreduced_kind to be of type `jax\\.sharding\\.UnreducedKind` but got (.+?)","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/partition_spec.py","lineNumber":38,"sourceCode":"from jax._src.lib import _jax\n\nAxisName = Any\n\ndef _check(partitions, unreduced, reduced, unreduced_kind):\n  if None in unreduced:\n    raise ValueError(\n        \"unreduced cannot contain None. All elements in unreduced should refer\"\n        \" to the mesh axes.\")\n  if None in reduced:\n    raise ValueError(\n        \"reduced cannot contain None. All elements in reduced should refer\"\n        \" to the mesh axes.\")\n  if unreduced & reduced:\n    raise ValueError(\n        \"`unreduced` and `reduced` argument to PartitionSpec cannot overlap. \"\n        f\"Got unreduced: {unreduced} and reduced: {reduced}\")\n  if unreduced_kind is not None and not isinstance(unreduced_kind, UnreducedKind):\n      raise TypeError(\n          \"Expected unreduced_kind to be of type `jax.sharding.UnreducedKind`\"\n          f\" but got {type(unreduced_kind)}\")\n  if not unreduced and unreduced_kind is not None:\n    raise ValueError(\n        \"`unreduced_kind` should be `None` when `unreduced` is an empty set.\"\n        f\" Got {unreduced_kind=} and {unreduced=}\")\n\n  for partition in partitions:\n    partition = partition if isinstance(partition, tuple) else (partition,)\n    for p in partition:\n      if p in unreduced:\n        raise ValueError(\n            \"partitions cannot overlap with unreduced axes passed to\"\n            f\" PartitionSpec. Got partitions: {partitions} and unreduced axes:\"\n            f\" {unreduced}\")\n      if p in reduced:\n        raise ValueError(\n            \"partitions cannot overlap with reduced axes passed to\"","sourceCodeStart":20,"sourceCodeEnd":56,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/partition_spec.py#L20-L56","documentation":"The optional unreduced_kind argument of the new-style PartitionSpec must be a jax.sharding.UnreducedKind instance (or None). Passing any other type — string, int, custom class — raises TypeError.","triggerScenarios":"PartitionSpec(..., unreduced=('data',), unreduced_kind='partial') or unreduced_kind=0.","commonSituations":"Passing a string flag from config (e.g. 'optimistic'/'partial') instead of the enum-like UnreducedKind object.","solutions":["Pass jax.sharding.UnreducedKind values (import from jax.sharding) or omit the argument","Check type(unreduced_kind) is UnreducedKind before constructing"],"exampleFix":"# before\nPartitionSpec(ps, unreduced=('d',), unreduced_kind='partial')\n# after\nfrom jax.sharding import UnreducedKind\nPartitionSpec(ps, unreduced=('d',), unreduced_kind=UnreducedKind.PARTIAL)","handlingStrategy":"type-guard","validationCode":"from jax.sharding import UnreducedKind\nassert unreduced_kind is None or isinstance(unreduced_kind, UnreducedKind)","typeGuard":"def valid_kind(k): return k is None or isinstance(k, UnreducedKind)","tryCatchPattern":null,"preventionTips":["Pass UnreducedKind members, not strings"],"tags":["jax","sharding","partition-spec","type-error"],"backgroundTag":"invalid-argument-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}