tensorflow/models · error · ValueError

The number of logical devices

Error message

The number of logical devices %d is not supported. Supported numbers are 1, 2, 4, 8, 16

What it means

Constraint guard in get_computation_shape_for_model_parallelism: spatial partitioning only has hard-coded computation shapes for 1, 2, 4, 8, or 16 logical devices (the product of input_partition_dims); any other device count has no valid 4-D computation shape, so the mapping refuses it.

Solutions

  1. Set the number of logical devices to one of 1, 2, 4, 8, or 16.
  2. Adjust the TPU topology or device partition config to a supported count.
Defensive patterns

Strategy: validation

When it happens

Trigger: Thrown at official/vision/train_spatial_partitioning.py:61 when the library encounters an invalid state.

Common situations: See trigger scenarios.

AI-assisted analysis of tensorflow/models@e006f5f0d5 (2026-08-24). Data as JSON: /api/errors/6a0187dd8aba874f. Report an issue: GitHub.

Appendix: source

Thrown at official/vision/train_spatial_partitioning.py:61

  Returns:
    A list of integers specifying the computation shape.

  Raises:
    ValueError: If the number of logical devices is not supported.
  """
  num_logical_devices = np.prod(input_partition_dims)
  if num_logical_devices == 1:
    return [1, 1, 1, 1]
  elif num_logical_devices == 2:
    return [1, 1, 1, 2]
  elif num_logical_devices == 4:
    return [1, 2, 1, 2]
  elif num_logical_devices == 8:
    return [2, 2, 1, 2]
  elif num_logical_devices == 16:
    return [4, 2, 1, 2]
  else:
    raise ValueError(
        'The number of logical devices %d is not supported. Supported numbers '
        'are 1, 2, 4, 8, 16' % num_logical_devices)


def create_distribution_strategy(distribution_strategy,
                                 tpu_address,
                                 input_partition_dims=None,
                                 num_gpus=None):
  """Creates distribution strategy to use for computation."""

  if input_partition_dims is not None:
    if distribution_strategy != 'tpu':
      raise ValueError('Spatial partitioning is only supported '
                       'for TPUStrategy.')

    # When `input_partition_dims` is specified create custom TPUStrategy
    # instance with computation shape for model parallelism.
    resolver = tf.distribute.cluster_resolver.TPUClusterResolver(

View on GitHub (pinned to e006f5f0d5)