TensorFlow
Split Op
input type error
type mismatch
debugging

TensorFlow - Input 'split_dim' of 'Split' Op has type float32 that does not match expected type of int32

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

This TensorFlow error means the axis argument passed to a split operation is a floating-point value when TensorFlow expects an integer axis. In older traces the parameter may appear as split_dim; in current APIs it is usually called axis, but the root problem is the same: the dimension index must be an integer type.

What TensorFlow expects

tf.split divides a tensor along a specified axis. That axis must be an int or a scalar integer tensor such as int32.

python
1import tensorflow as tf
2
3x = tf.constant([[1, 2], [3, 4], [5, 6]])
4parts = tf.split(x, num_or_size_splits=2, axis=1)
5
6for part in parts:
7    print(part)

This works because axis=1 is an integer. The error appears when the axis comes in as something like 1.0, tf.float32, or a tensor that was created from floating-point data.

A common failing example

Here is the kind of code that triggers the error:

python
1import tensorflow as tf
2
3x = tf.constant([[1, 2], [3, 4]])
4axis = tf.constant(1.0, dtype=tf.float32)
5
6tf.split(x, num_or_size_splits=2, axis=axis)

The shape dimension is conceptually an index, and indexes are integers. TensorFlow refuses the float because splitting along "dimension 1.0" is not a valid graph type.

Fix the axis type

If the axis is meant to be constant, pass a normal integer:

python
axis = 1
parts = tf.split(x, num_or_size_splits=2, axis=axis)

If the axis comes from tensor-based logic, make it an integer tensor:

python
axis = tf.constant(1, dtype=tf.int32)
parts = tf.split(x, num_or_size_splits=2, axis=axis)

If the value arrives as a float but logically represents a whole-number axis, cast it explicitly:

python
axis = tf.constant(1.0, dtype=tf.float32)
axis = tf.cast(axis, tf.int32)
parts = tf.split(x, num_or_size_splits=2, axis=axis)

That cast is appropriate only if the value is truly intended to be an integer dimension.

The same rule applies to related APIs such as tf.unstack and some shape-manipulation ops: tensor data may be floating point, but shape positions and axes should stay integer typed.

Trace the value back to its source

In real code, the axis is often not hard-coded. It may come from:

  • parsed configuration
  • NumPy arrays converted with the wrong dtype
  • placeholder tensors
  • arithmetic that accidentally produces floats

A quick debug pattern is:

python
print(axis)
print(getattr(axis, "dtype", type(axis)))

If you are inside TensorFlow graph code, also inspect intermediate tensors before they reach tf.split. The best fix is usually to keep dimension values as integers from the start instead of casting them late.

Do not confuse axis with num_or_size_splits

Another source of confusion is that tf.split has two important parameters:

  • 'num_or_size_splits'
  • 'axis'

The error in this article is about axis, not about the split sizes. If num_or_size_splits is wrong, TensorFlow raises a different error, usually about divisibility or inconsistent sizes rather than about int32 versus float32.

Common Pitfalls

The biggest pitfall is letting numeric metadata flow through generic float pipelines. Shape indices, split axes, and array positions should stay as integers even if the model data itself is floating point.

Another issue is assuming Python numeric literals and tensor dtypes behave the same way. 1 is fine as an axis, but tf.constant(1.0) is not.

It is also easy to patch the immediate error with a cast and miss the upstream bug that produced a float axis in the first place. If the value was never supposed to be float, fix it earlier.

Finally, remember that TensorFlow error messages may use legacy names such as split_dim depending on the API layer or version. Focus on the meaning of the parameter, not just the exact label in the traceback.

Summary

  • 'tf.split expects the axis to be an integer, not float32.'
  • Pass axis as a Python int or an integer tensor such as tf.int32.
  • Cast only when the float value is logically an integer dimension.
  • Trace the source of the axis value if the wrong dtype appears repeatedly.
  • Do not confuse axis-type errors with split-size errors.

Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Basis Lab
No course covers this one yet

Basis Lab writes one for you from a sentence about what you want to be able to do, then teaches it and checks you understood. Your first course is free.

Build my course
ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design