from tensorflow.python.framework import tensor_shape

    returned_shape = []
    for i, dim in enumerate(input_shape.dims):
      if i != dimension:
        returned_shape.append(dim)
    return [tensor_shape.TensorShape(returned_shape)]
  else:
    raise ValueError(
        "dimension (%d) must be in the range [0, %d), where %d is the number "
