tensorflow/models · error · ValueError

Spatial partitioning is only supported for TPUStrategy.

Error message

Spatial partitioning is only supported for TPUStrategy.

What it means

Error "Spatial partitioning is only supported for TPUStrategy." thrown in tensorflow/models.

Source

Thrown at official/vision/train_spatial_partitioning.py:74

  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(
        tpu=tpu_address)
    if tpu_address not in ('', 'local'):
      tf.config.experimental_connect_to_cluster(resolver)

    topology = tf.tpu.experimental.initialize_tpu_system(resolver)
    num_replicas = resolver.get_tpu_system_metadata().num_cores // np.prod(
        input_partition_dims)
    device_assignment = tf.tpu.experimental.DeviceAssignment.build(
        topology,
        num_replicas=num_replicas,
        computation_shape=input_partition_dims)
    return tf.distribute.TPUStrategy(
        resolver, experimental_device_assignment=device_assignment)

View on GitHub (pinned to e006f5f0d5)

Solutions

1. Run under TPUStrategy when using spatial partitioning.
2. Disable spatial partitioning when training on GPU/CPU.

When it happens

Trigger: Thrown at official/vision/train_spatial_partitioning.py:74 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/90fbd6a7beafaf18. Report an issue: GitHub.