{"record":{"id":"8d7b804380f942d0","repo":"jax-ml/jax","slug":"name-was-requested-to-map-a-value-of-non-array-t","errorCode":null,"errorMessage":"{name} was requested to map a value of non-array type {core.typeof(x)} along axis {axis}, but non-array types can't be mapped along an integer axis. Instead pass a mapping spec (a MappingSpec instance) as this argument's in_axes entry, and pass axis_size explicitly.","messagePattern":"(.+?) was requested to map a value of non-array type (.+?) along axis (.+?), but non-array types can't be mapped along an integer axis\\. Instead pass a mapping spec \\(a MappingSpec instance\\) as this argument's in_axes entry, and pass axis_size explicitly\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/api.py","lineNumber":1316,"sourceCode":"    if axis_size is not None:\n      return axis_size\n    args, kwargs = tree_unflatten(tree, vals)\n    raise ValueError(\n        f\"{name} wrapped function must be passed at least one argument \"\n        \"containing an array or axis_size must be specified, got empty \"\n        f\"*args={args} and **kwargs={kwargs}\"\n    )\n\n  def _get_axis_size(name: str, x, axis: int) -> core.AxisSize | None:\n    shape: tuple[core.AxisSize, ...] = ()\n    try:\n      shape = np.shape(x)\n      return shape[axis]\n    except (IndexError, TypeError) as e:\n      if not core.valid_jaxtype(x) or not isinstance(axis, int):\n        return None  # Suppress the check for custom vmappable types.\n      if core.typeof(x).is_high:\n        raise ValueError(\n            f\"{name} was requested to map a value of non-array type \"\n            f\"{core.typeof(x)} along axis {axis}, but non-array types can't \"\n            \"be mapped along an integer axis. Instead pass a mapping spec (a \"\n            \"MappingSpec instance) as this argument's in_axes entry, and \"\n            \"pass axis_size explicitly.\") from None\n      min_rank = axis + 1 if axis >= 0 else -axis\n      # TODO(mattjj): better error message here\n      raise ValueError(\n          f\"{name} was requested to map its argument along axis {axis}, \"\n          f\"which implies that its rank should be at least {min_rank}, \"\n          f\"but is only {len(shape)} (its shape is {shape})\") from e\n\n  all_mapped_sizes = [\n    None if d is None else _get_axis_size(name, x, d)\n    for x, d in zip(vals, dims)\n  ]\n  all_sizes = [s for s in all_mapped_sizes if s is not None]\n  if axis_size is not None:","sourceCodeStart":1298,"sourceCodeEnd":1334,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/api.py#L1298-L1334","documentation":"Raised when vmap/pmap/linearize tries to map a non-array JAX type (e.g. a string, dict, or custom high-level type) along an integer in_axes entry. Non-array pytree nodes cannot have an integer axis mapped; JAX requires a MappingSpec and an explicit axis_size instead.","triggerScenarios":"jax.vmap(f)(x, {'a': 1}) with in_axes=0 where the dict argument is a non-array leaf; mapping a custom type with a plain integer in_axes.","commonSituations":"Passing dicts of metadata or custom vmappable objects through vmap with default in_axes=0; new JAX versions that enforce MappingSpec for high-level types.","solutions":["Set in_axes=None for the non-array argument","Pass a MappingSpec instance for that argument's in_axes entry and axis_size explicitly","Convert the metadata to static/closure state or a JAX array"],"exampleFix":"// before\njax.vmap(f, in_axes=(0, 0))(x, meta_dict)\n// after\njax.vmap(f, in_axes=(0, None))(x, meta_dict)","handlingStrategy":"type-guard","validationCode":"for leaf, ax in zip(tree_leaves(args), tree_leaves(in_axes)):\n    if isinstance(ax, int) and not (hasattr(leaf, 'ndim') or core.valid_jaxtype(leaf) and not core.typeof(leaf).is_high):\n        raise TypeError('non-array leaf needs in_axes=None or MappingSpec')","typeGuard":"def is_mappable_array(x): return hasattr(x, 'shape') and hasattr(x, 'dtype')","tryCatchPattern":null,"preventionTips":["Set in_axes=None for metadata/dict leaves","Use MappingSpec + axis_size for custom mapped types","Validate leaves with hasattr(x, 'ndim') before vmap"],"tags":["jax","vmap","in-axes","type-error"],"backgroundTag":"invalid-argument-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}