a39084b995d9d6d88300632c54895488985e5cfd,onnx_tf/backend.py,TensorflowBackend,handle_flatten,#Any#Any#Any#,237

Before Change

    y_direction = 1;
    for i in range(axis):
      x_direction = x_direction * tensor.shape[i]
    for i in range(axis, tensor.shape.length):
      y_direction = y_direction * tensor.shape[i]
    shape = tf.constant([x_direction, y_direction])
    return [tf.reshape(tensor, shape)]

  @classmethod

After Change

  def handle_flatten(cls, node, input_dict):
    tensor = input_dict[node.inputs[0]]
    axis = node.attrs["axis"] if "axis" in node.attrs.keys() else 1
    shape = tf.shape(tensor)
    split0, split1 = tf.split(shape, [axis, tf.size(shape) - axis])
    split0 = tf.reduce_prod(split0)
    split1 = tf.reduce_prod(split1)
    output_shape = tf.stack([split0, split1])
    return [tf.reshape(tensor, output_shape)]

Instances