[nnx] improving performance of _graph_flatten and `_graph_unflatten…

[erfanzar]

Optimize _graph_flatten and _graph_unflatten functions

What does this PR do?

This PR optimizes the performance of _graph_flatten and _graph_unflatten functions in graph.py. The main changes include:

• Using isinstance() instead of type() for type checking, which is more efficient
• Simplified code structure for better readability and performance
• Improved overall execution time
• Early Return for Cached Nodes
• Efficient Index Assignment
• Minimized Recursive Calls
• Reduced Redundant Operations (Avoided recomputing values that can be stored and reused.)

Performance improvements:
Before:

• Total time: 15.339856624603271s
• Time per step: 15339.86µs
• Time per layer: 153.40µs
• _graph_flatten: 1.798s
• _graph_unflatten: 0.4857s

After:

• Total time: 14.821305990219116s
• Time per step: 14821.31µs
• Time per layer: 148.21µs
• _graph_flatten: 1.734s
• _graph_unflatten: 0.4772s

These optimizations result in approximately 3.4% improvement in overall execution time.

Fixes # (issue)

Checklist

• This PR fixes a minor issue (e.g.: typo or small bug) or improves the docs.
• This change is discussed in a Github issue/discussion.
• The documentation and docstrings adhere to the documentation guidelines.
• This change includes necessary high-coverage tests.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[cgarciae]

why remove the annotations?