{"record":{"id":"e05d446b04264eac","repo":"jax-ml/jax","slug":"lax-associative-scan-fn-argument-should-be-callab","errorCode":null,"errorMessage":"lax.associative_scan: fn argument should be callable.","messagePattern":"lax\\.associative_scan: fn argument should be callable\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/control_flow/loops.py","lineNumber":2877,"sourceCode":"\n  Example 2: partial products of an array of matrices\n\n  >>> mats = jax.random.uniform(jax.random.key(0), (4, 2, 2))\n  >>> partial_prods = lax.associative_scan(jnp.matmul, mats)\n  >>> partial_prods.shape\n  (4, 2, 2)\n\n  Example 3: reversed partial sums of an array of numbers\n\n  >>> lax.associative_scan(jnp.add, jnp.arange(0, 4), reverse=True)\n  Array([6, 6, 5, 3], dtype=int32)\n\n  .. [BLE1990] Blelloch, Guy E. 1990. \"Prefix Sums and Their Applications.\",\n    Technical Report CMU-CS-90-190, School of Computer Science, Carnegie Mellon\n    University.\n  \"\"\"\n  if not callable(fn):\n    raise TypeError(\"lax.associative_scan: fn argument should be callable.\")\n  elems_flat, tree = tree_flatten(elems)\n\n  if reverse:\n    elems_flat = [lax.rev(elem, [axis]) for elem in elems_flat]\n\n  def combine(a_flat, b_flat):\n    # Lower `fn` to operate on flattened sequences of elems.\n    a = tree_unflatten(tree, a_flat)\n    b = tree_unflatten(tree, b_flat)\n    c = fn(a, b)\n    c_flat, _ = tree_flatten(c)\n    return c_flat\n\n  # Check that all inputs have a consistent leading dimension `num_elems`.\n  axis = util.canonicalize_axis(axis, elems_flat[0].ndim)\n\n  if not core.is_constant_dim(elems_flat[0].shape[axis]):\n    raise NotImplementedError(\"associative scan over axis \"","sourceCodeStart":2859,"sourceCodeEnd":2895,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/control_flow/loops.py#L2859-L2895","documentation":"jax.lax.associative_scan requires its first argument to be a Python callable (the binary associative combine function). Before flattening the inputs, the implementation checks `callable(fn)` and raises TypeError immediately if it fails, since there is no meaningful way to perform a scan without a combine operation.","triggerScenarios":"Calling lax.associative_scan(fn, elems) where fn is not callable, e.g. passing a string like '+', an operator symbol, a jnp array, or forgetting the function argument entirely and passing only arrays.","commonSituations":"Developers coming from other libraries where reductions accept string op names (e.g. numpy/jnp reductions or torch); passing jnp.add result instead of jnp.add; refactoring code and dropping the fn argument.","solutions":["Pass a callable as the first argument, e.g. jnp.add or a lambda: lax.associative_scan(jnp.add, elems)","If you wanted a string op, wrap it: lax.associative_scan(lambda a, b: a + b, elems)","Check argument order — signature is associative_scan(fn, elems, reverse=False)"],"exampleFix":"// before\nlax.associative_scan('+', elems)\n// after\nlax.associative_scan(jnp.add, elems)","handlingStrategy":"type-guard","validationCode":"assert callable(fn), 'fn must be callable'","typeGuard":"def is_scan_fn(fn) -> bool:\n    return callable(fn)","tryCatchPattern":null,"preventionTips":["Pass jnp.add or an explicit lambda; associative_scan never accepts string ops","Lint call sites where the first positional arg is a str literal"],"tags":["jax","associative-scan","type-error","callable"],"backgroundTag":"non-callable-argument","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}