{"record":{"id":"5f657c3968357607","repo":"jax-ml/jax","slug":"encountered-unexpected-shape-dimension-d","errorCode":null,"errorMessage":"Encountered unexpected shape dimension {d}","messagePattern":"Encountered unexpected shape dimension (.+?)","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/export/shape_poly.py","lineNumber":1284,"sourceCode":"  opt_einsum.contract_path to parse the specification.\n  \"\"\"\n\n  # Replace the polymorphic shapes with some concrete shapes for calling\n  # into opt_einsum.contract_path, because the latter wants to compute the\n  # sizes of operands and intermediate results.\n  fake_ops = []\n  for operand in operands:\n    # We replace only array operands\n    if not hasattr(operand, \"dtype\"):\n      fake_ops.append(operand)\n    else:\n      shape = np.shape(operand)\n      def fake_dim(d):\n        if core.is_constant_dim(d):\n          return d\n        else:\n          if not isinstance(d, _DimExpr):\n            raise TypeError(f\"Encountered unexpected shape dimension {d}\")\n          # It is Ok to replace all polynomials with the same value. We may miss\n          # here some errors due to non-equal dimensions, but we catch them\n          # later.\n          return 8\n      fake_ops.append(api.ShapeDtypeStruct(tuple(map(fake_dim, shape)),\n                                           operand.dtype))\n\n  contract_fake_ops, contractions = opt_einsum.contract_path(*fake_ops,\n                                                             **kwargs)\n  contract_operands = []\n  for operand in contract_fake_ops:\n    idx = tuple(i for i, fake_op in enumerate(fake_ops) if operand is fake_op)\n    assert len(idx) == 1\n    contract_operands.append(operands[idx[0]])\n  return contract_operands, contractions\n\n# To implement shape-constraint checking we use a shape assertion primitive.\n#    shape_assertion_p.bind(assert_what: bool, *error_message_inputs,","sourceCodeStart":1266,"sourceCodeEnd":1302,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/export/shape_poly.py#L1266-L1302","documentation":"During export, JAX substitutes fake constant dims for symbolic ones when preparing operands for underlying lowering. If a non-constant dimension is not a _DimExpr (e.g. a numpy object, a string, or a dim from a foreign tracer), it cannot be faked and raises TypeError.","triggerScenarios":"An operand's shape contains a non-int, non-_DimExpr dimension object — e.g. a Tracer from another transformation, np.str_, or a custom dim class — while exporting with polymorphic shapes.","commonSituations":"Composing jax.export with other transforms producing exotic dim objects; passing arrays whose shapes were built by external libs inserting object dims; JAX version incompatibilities in the dim representation.","solutions":["Ensure all non-constant dims in operand shapes are genuine symbolic dims created by the current export's shape polymorphism","Materialize shapes to plain ints / re-create inputs as jnp arrays before export","Align jax/jaxlib versions so dim objects are the internal _DimExpr type"],"exampleFix":"# before\nx = some_external_array  # shape contains object dims\nexp = jax.export.export(shapes('n,'))(f)\n# after\nx = jnp.asarray(x)  # normalize to a JAX array with int/symbolic dims\nexp = jax.export.export(shapes('n,'))(f)","handlingStrategy":"validation","validationCode":"def clean_shape(a):\n    return all(isinstance(d, int) or is_symbolic_dim(d) for d in np.shape(a))\nassert clean_shape(x), 'operand has non-int/non-symbolic dims'","typeGuard":"def exportable_operand(a) -> bool:\n    return isinstance(a, jax.Array) or all(isinstance(d, (int,)) or hasattr(d, '_factors') for d in np.shape(a))","tryCatchPattern":"try:\n    exp = jax.export.export(shapes)(fn)\nexcept TypeError as e:\n    if 'unexpected shape dimension' in str(e): inputs = [jnp.asarray(i) for i in inputs]; retry()","preventionTips":["Convert external arrays to jnp arrays before export","Keep jax/jaxlib versions aligned so dim types match"],"tags":["jax","export","shape-polymorphism","type-error"],"backgroundTag":"invalid-shape-dimension-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}