tensorflow/models · error · ValueError

Spatial partitioning is only supported for TPUStrategy.

Error message

Spatial partitioning is only supported for TPUStrategy.

What it means

Capability guard in create_distribution_strategy: spatial partitioning requires partitioned vars/ops only TPUStrategy provides; when input_partition_dims is set but the resolved strategy is not a TPUStrategy, the function refuses to build a setup where partitioned inputs could not be mapped to devices.

Solutions

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

Strategy: validation

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.

Appendix: 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)