jax-ml/jax · error · ValueError

No VJP is available

Error message

No VJP is available

What it means

Exported.vjp() returns the VJP of a serialized/exported function. If the Exported was loaded from an external serialization format that did not include VJP data, _get_vjp is None and this ValueError is raised.

Source

Thrown at jax/_src/export/_export.py:295

    See documentation for in_shardings_jax.
    """
    return tuple(
      _get_named_sharding(named_sharding, mesh)
      for named_sharding in self._out_named_shardings)

  def has_vjp(self) -> bool:
    """Returns if this Exported supports VJP."""
    return self._get_vjp is not None

  def vjp(self) -> Exported:
    """Gets the exported VJP.

    Returns None if not available, which can happen if the Exported has been
    loaded from an external format without a VJP.
    """
    if self._get_vjp is None:
      raise ValueError("No VJP is available")
    return self._get_vjp(self)

  def serialize(self,
                vjp_order: int = 0) -> bytearray:
    """Serializes an Exported.

    Args:
      vjp_order: The maximum vjp order to include. E.g., the value 2 means that we
        serialize the primal functions and two orders of the ``vjp`` function. This
        should allow 2nd order reverse mode differentiation of the deserialized
        function. i.e., ``jax.grad(jax.grad(f))``.
    """
    # Lazy load the serialization module, since flatbuffers is an optional
    # dependency.
    from jax._src.export.serialization import serialize
    return serialize(self, vjp_order=vjp_order)

  def call(self, *args, **kwargs):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Check availability first: if exported.vjp_available (or inspect _get_vjp) before calling vjp()
2. Re-export from the original JAX program with VJP included (serialize with the vjp captured)
3. For loaded exports, reconstruct gradients via a differentiating wrapper (jax.vjp) around the loaded callable instead

Example fix

# before
exp = load_exported(path)
vjp_fn = exp.vjp()  # ValueError

# after
exp = load_exported(path)
vjp_fn = jax.vjp(exp.call) if not exp_has_vjp(exp) else exp.vjp()

has_vjp = getattr(exported, '_get_vjp', None) is not None
vjp_fn = exported.vjp() if has_vjp else jax.vjp(exported.call)

def exported_has_vjp(exp) -> bool:
    return getattr(exp, '_get_vjp', None) is not None

try:
    vjp_fn = exported.vjp()
except ValueError:
    vjp_fn = jax.vjp(exported.call)  # differentiate around the loaded callable

When it happens

Trigger: Loading an Exported from a StableHLO/external file (not created in-process with a VJP) and then calling .vjp(); deserializing a model saved for inference-only and requesting gradients; calling vjp() with vjp_order larger than what was exported.

Common situations: Serving/deployment pipelines that load exported models and later attempt fine-tuning or gradient-based analysis; mixing in-process exports (which carry VJPs) with file-loaded ones.

Related errors

• Duplicate serialization registration for type `{nodetype}`.
• Duplicate serialization registration for serialized_name `{s
• serializing DisabledSafetyCheck: {check}
• reduce_axes argument to vjp is deprecated
• unexpected JAX type (e.g. shape/dtype) for gradient ref pass

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/c7a2f56be3f85ad5. Report an issue: GitHub.