Remove axis argument - it will always be 0 when called by predict
@@ -3258,14 +3258,14 @@ def concat(tensors, axis=0):
return tf.concat(tensors, axis=axis)
-def potentially_variable_concat(tensors, axis=0):
+def potentially_variable_concat(tensors):
"""Concats `tensor`s along `axis`. When non-batch dimensions are variable, a RaggedTensor is returned instead."""
if isinstance(tensors[0], tf.SparseTensor):
- return tf.sparse.concat(axis=axis, sp_inputs=tensors)
+ return tf.sparse.concat(axis=0, sp_inputs=tensors)
non_batch_shapes = tf.stack([tf.shape(tensor)[1:] for tensor in tensors])
constant_dims = tf.math.reduce_all(non_batch_shapes == non_batch_shapes[:1], axis=0)
if tf.math.reduce_all(constant_dims).numpy().item(): # All non-batch dimensions are constant
- return tf.concat(tensors, axis=axis)
+ return tf.concat(tensors, axis=0)
# First, identify constant inner dimensions by finding the rightmost dimension that is not constant
constant_inner_dimensions = constant_dims.numpy().tolist()[::-1].index(False)
# If there are constant inner dimensions, define a constant inner shape
