from tensorflow.python.framework import tensor_shape

  else:
    for i, dim in enumerate(input_shape.dims):
      if i not in reduction_indices:
        returned_dims.append(dim)
  return [tensor_shape.TensorShape(returned_dims)]


@ops.RegisterShape("SegmentMax")
