keras-team/keras · error · ValueError

Argument `num_segments` must be set when using the JAX backe

Error message

Argument `num_segments` must be set when using the JAX backend. Received: num_segments=None

What it means

Error "Argument `num_segments` must be set when using the JAX backend. Received: num_segments=None" thrown in keras-team/keras.

Source

Thrown at keras/src/backend/jax/ops/math.py:13

import math

import jax
import jax.numpy as jnp

from keras.src.backend.common.variables import standardize_dtype
from keras.src.backend.jax.ops.core import convert_to_tensor
from keras.src.utils.module_utils import scipy


def segment_sum(data, segment_ids, num_segments=None, sorted=False):
    if num_segments is None:
        raise ValueError(
            "Argument `num_segments` must be set when using the JAX backend. "
            "Received: num_segments=None"
        )
    return jax.ops.segment_sum(
        data, segment_ids, num_segments, indices_are_sorted=sorted
    )


def segment_max(data, segment_ids, num_segments=None, sorted=False):
    if num_segments is None:
        raise ValueError(
            "Argument `num_segments` must be set when using the JAX backend. "
            "Received: num_segments=None"
        )
    return jax.ops.segment_max(
        data, segment_ids, num_segments, indices_are_sorted=sorted
    )

View on GitHub (pinned to 7a34a03db6)

When it happens

Trigger: Thrown at keras/src/backend/jax/ops/math.py:13 when the library encounters an invalid state.

Common situations: See trigger scenarios.

AI-assisted analysis of keras-team/keras@7a34a03db6 (2026-08-25). Data as JSON: /api/errors/06a461456d111b71. Report an issue: GitHub.