Move the python _infer_shape_and_flatten function to C++

[kellen-sun]

Pretty simple like the title says. There's no good reason to keep it in Python, especially as it's not interface relevant.
Can just have the function and pybind it; refer to bindings.cpp and bindings_impl.cpp for examples.