Skip to content

Commit

Permalink
Update: ivy.compile( ) updated to ivy.trace_graph( ) (#22764)
Browse files Browse the repository at this point in the history
Co-authored-by: ivy-branch <[email protected]>
  • Loading branch information
HaiderSultanArc and ivy-branch authored Oct 2, 2023
1 parent c01ad72 commit 39671b3
Show file tree
Hide file tree
Showing 18 changed files with 145 additions and 145 deletions.
10 changes: 5 additions & 5 deletions README.md
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
> 🚀 We are granting pilot access to **Ivy\'s Compiler and Transpiler**
> 🚀 We are granting pilot access to **Ivy\'s Tracer and Transpiler**
> to some users, [join the waitlist](https://console.unify.ai/) if you
> want to test them out!
Expand Down Expand Up @@ -131,8 +131,8 @@ deploy systems. Feel free to head over to the docs for the full API
reference, but the functions you\'d most likely want to use are:

``` python
# Compiles a function into an efficient fully-functional graph, removing all wrapping and redundant code
ivy.compile()
# Traces an efficient fully-functional graph from a function, removing all wrapping and redundant code
ivy.trace_graph()

# Converts framework-specific code to a different framework
ivy.transpile()
Expand All @@ -142,8 +142,8 @@ ivy.unify()
```

These functions can be used eagerly or lazily. If you pass the necessary
arguments for function tracing, the compilation/transpilation step will
happen instantly (eagerly). Otherwise, the compilation/transpilation
arguments for function tracing, the tracing/transpilation step will
happen instantly (eagerly). Otherwise, the tracing/transpilation
will happen only when the returned function is first invoked.

``` python
Expand Down
6 changes: 3 additions & 3 deletions docs/overview/contributing/error_handling.rst
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ This section, "Error Handling" aims to assist you in navigating through some com
E with_out=False,
E instance_method=False,
E test_gradients=False,
E test_compile=None,
E test_trace=None,
E as_variable=[False],
E native_arrays=[False],
E container=[False],
Expand Down Expand Up @@ -65,7 +65,7 @@ This section, "Error Handling" aims to assist you in navigating through some com
E with_out=False,
E instance_method=False,
E test_gradients=True,
E test_compile=None,
E test_trace=None,
E as_variable=[False],
E native_arrays=[False],
E container=[False],
Expand Down Expand Up @@ -129,7 +129,7 @@ This section, "Error Handling" aims to assist you in navigating through some com
E with_out=False,
E instance_method=False,
E test_gradients=False,
E test_compile=None,
E test_trace=None,
E as_variable=[False],
E native_arrays=[False],
E container=[False],
Expand Down
4 changes: 2 additions & 2 deletions docs/overview/deep_dive/containers.rst
Original file line number Diff line number Diff line change
Expand Up @@ -252,8 +252,8 @@ There may be some compositional functions which are not implicitly nestable for
One such example is the :func:`ivy.linear` function which is not implicitly nestable despite being compositional. This is because of the use of special functions like :func:`__len__` and :func:`__list__` which, among other functions, are not nestable and can't be made nestable.
But we should try to avoid this, in order to make the flow of computation as intuitive to the user as possible.

When compiling the code, the computation graph is **identical** in either case, and there will be no implications on performance whatsoever.
The implicit nestable solution may be slightly less efficient in eager mode, as the leaves of the container are traversed multiple times rather than once, but if performance is of concern then the code should always be compiled in any case.
When tracing the code, the computation graph is **identical** in either case, and there will be no implications on performance whatsoever.
The implicit nestable solution may be slightly less efficient in eager mode, as the leaves of the container are traversed multiple times rather than once, but if performance is of concern then the code should always be traced in any case.
The distinction is only really relevant when stepping through and debugging with eager mode execution, and for the reasons outlined above, the preference is to keep compositional functions implicitly nestable where possible.

**Shared Nested Structure**
Expand Down
6 changes: 3 additions & 3 deletions docs/overview/deep_dive/ivy_frontends.rst
Original file line number Diff line number Diff line change
Expand Up @@ -92,12 +92,12 @@ The former set of functions map very closely to the API for the Accelerated Line
The latter set of functions map very closely to NumPy's well known API.
In general, all functions in the :mod:`jax.numpy` namespace are themselves implemented as a composition of the lower-level functions in the :mod:`jax.lax` namespace.

When transpiling between frameworks, the first step is to compile the computation graph into low level python functions for the source framework using Ivy's graph compiler, before then replacing these nodes with the associated functions in Ivy's frontend API.
When transpiling between frameworks, the first step is to trace a computation graph of low level python functions for the source framework using Ivy's tracer, before then replacing these nodes with the associated functions in Ivy's frontend API.
Given that all jax code can be decomposed into :mod:`jax.lax` function calls, when transpiling JAX code it should always be possible to express the computation graph as a composition of only :mod:`jax.lax` functions.
Therefore, arguably these are the *only* functions we should need to implement in the JAX frontend.
However, in general we wish to be able to compile a graph in the backend framework with varying levels of dynamicism.
However, in general we wish to be able to trace a graph in the backend framework with varying levels of dynamicism.
A graph of only :mod:`jax.lax` functions chained together in general is more *static* and less *dynamic* than a graph which chains :mod:`jax.numpy` functions together.
We wish to enable varying extents of dynamicism when compiling a graph with our graph compiler, and therefore we also implement the functions in the :mod:`jax.numpy` namespace in our frontend API for JAX.
We wish to enable varying extents of dynamicism when creating a graph with our tracer, and therefore we also implement the functions in the :mod:`jax.numpy` namespace in our frontend API for JAX.

Thus, both :mod:`lax` and :mod:`numpy` modules are created in the JAX frontend API.
We start with the function :func:`lax.add` as an example.
Expand Down
8 changes: 4 additions & 4 deletions docs/overview/deep_dive/superset_behaviour.rst
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,7 @@ We've already explained that we should not duplicate arguments in the Ivy functi
Does this mean, provided that the proposed argument is not a duplicate, that we should always add this backend-specific argument to the Ivy function?
The answer is **no**.
When determining the superset, we are only concerned with the pure **mathematics** of the function, and nothing else.
For example, the :code:`name` argument is common to many TensorFlow functions, such as `tf.concat <https://www.tensorflow.org/api_docs/python/tf/concat>`_, and is used for uniquely identifying parts of the compiled computation graph during logging and debugging.
For example, the :code:`name` argument is common to many TensorFlow functions, such as `tf.concat <https://www.tensorflow.org/api_docs/python/tf/concat>`_, and is used for uniquely identifying parts of the traced computation graph during logging and debugging.
This has nothing to do with the mathematics of the function, and so is *not* included in the superset considerations when implementing Ivy functions.
Similarly, in NumPy the argument :code:`subok` controls whether subclasses of the :class:`numpy.ndarray` class should be permitted, which is included in many functions, such as `numpy.ndarray.astype <https://numpy.org/doc/stable/reference/generated/numpy.ndarray.astype.html>`_.
Finally, in JAX the argument :code:`precision` is quite common, which controls the precision of the return values, as used in `jax.lax.conv <https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.conv.html>`_ for example.
Expand Down Expand Up @@ -129,8 +129,8 @@ The following would be a much better solution:
return res
You will notice that this implementation involves more lines of code, but this should not be confused with added complexity.
All Ivy code should be graph compiled for efficiency, and in this case all the :code:`if` and :code:`else` statements are removed, and all that remains is the backend functions which were executed.
This new implementation will be compiled to a graph of either one, three, four, or six functions depending on the values of :code:`beta` and :code:`threshold`, while the previous implementation would *always* compile to six functions.
All Ivy code should be traced for efficiency, and in this case all the :code:`if` and :code:`else` statements are removed, and all that remains is the backend functions which were executed.
This new implementation will be traced to a graph of either one, three, four, or six functions depending on the values of :code:`beta` and :code:`threshold`, while the previous implementation would *always* traces to six functions.

This does mean we do not adopt the default values used by PyTorch, but that's okay.
Implementing the superset does not mean adopting the same default values for arguments, it simply means equipping the Ivy function with the capabilities to execute the superset of behaviours.
Expand Down Expand Up @@ -167,7 +167,7 @@ With regards to both of these points, Ivy provides the generalized superset impl

However, as discussed above, :func:`np.logical_and` also supports the :code:`where` argument, which we opt to **not** support in Ivy.
This is because the behaviour can easily be created as a composition like so :code:`ivy.where(mask, ivy.logical_and(x, y), ivy.zeros_like(mask))`, and we prioritize the simplicity, clarity, and function uniqueness in Ivy's API in this case, which comes at the cost of reduced runtime efficiency for some functions when using a NumPy backend.
However, in future releases our automatic graph compilation and graph simplification processes will alleviate these minor inefficiencies entirely from the final computation graph, by fusing multiple operations into one at the API level where possible.
However, in future releases our automatic graph tracing and graph simplification processes will alleviate these minor inefficiencies entirely from the final computation graph, by fusing multiple operations into one at the API level where possible.

Maximizing Usage of Native Functionality
----------------------------------------
Expand Down
2 changes: 1 addition & 1 deletion docs/overview/design.rst
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ If that sounds like you, feel free to check out the `Deep Dive`_ section after y
| back-end functional APIs ✅
| Ivy functional API ✅
| Framework Handler ✅
| Ivy Compiler 🚧
| Ivy Tracer 🚧
|
| (b) `Ivy as a Transpiler <design/ivy_as_a_transpiler.rst>`_
| front-end functional APIs 🚧
Expand Down
40 changes: 20 additions & 20 deletions docs/overview/design/building_blocks.rst
Original file line number Diff line number Diff line change
Expand Up @@ -355,26 +355,26 @@ A good example is :func:`ivy.lstm_update`, as shown:
We *could* find and wrap the functional LSTM update methods for each backend framework which might bring a small performance improvement, but in this case there are no functional LSTM methods exposed in the official functional APIs of the backend frameworks, and therefore the functional LSTM code which does exist for the backends is much less stable and less reliable for wrapping into Ivy.
Generally, we have made decisions so that Ivy is as stable and scalable as possible, minimizing dependencies to backend framework code where possible with minimal sacrifices in performance.

Graph Compiler 🚧
Tracer 🚧
-----------------

“What about performance?” I hear you ask.
This is a great point to raise!

With the design as currently presented, there would be a small performance hit every time we call an Ivy function by virtue of the added Python wrapping.
One reason we created the graph compiler was to address this issue.
One reason we created the tracer was to address this issue.

The compiler takes in any Ivy function, backend function, or composition, and returns the computation graph using the backend functional API only.
The tracer takes in any Ivy function, backend function, or composition, and returns the computation graph using the backend functional API only.
The dependency graph for this process looks like this:

.. image:: https://github.com/unifyai/unifyai.github.io/blob/main/img/externally_linked/design/compiler_dependency_graph.png?raw=true
:align: center
:width: 75%

Let's look at a few examples, and observe the compiled graph of the Ivy code against the native backend code.
Let's look at a few examples, and observe the traced graph of the Ivy code against the native backend code.
First, let's set our desired backend as PyTorch.
When we compile the three functions below, despite the fact that each
has a different mix of Ivy and PyTorch code, they all compile to the same graph:
When we trace the three functions below, despite the fact that each
has a different mix of Ivy and PyTorch code, they all trace to the same graph:

+----------------------------------------+-----------------------------------------+-----------------------------------------+
|.. code-block:: python |.. code-block:: python |.. code-block:: python |
Expand All @@ -393,7 +393,7 @@ has a different mix of Ivy and PyTorch code, they all compile to the same graph:
| x = ivy.array([[1., 2., 3.]]) | x = torch.tensor([[1., 2., 3.]]) | x = ivy.array([[1., 2., 3.]]) |
| | | |
| # create graph | # create graph | # create graph |
| graph = ivy.compile_graph( | graph = ivy.compile_graph( | graph = ivy.compile_graph( |
| graph = ivy.trace_graph( | graph = ivy.trace_graph( | graph = ivy.trace_graph( |
| pure_ivy, x) | pure_torch, x) | mix, x) |
| | | |
| # call graph | # call graph | # call graph |
Expand All @@ -408,7 +408,7 @@ For all existing ML frameworks, the functional API is the backbone that underpin
This means that under the hood, any code can be expressed as a composition of ops in the functional API.
The same is true for Ivy.
Therefore, when compiling the graph with Ivy, any higher-level classes or extra code which does not directly contribute towards the computation graph is excluded.
For example, the following 3 pieces of code all compile to the exact same computation graph as shown:
For example, the following 3 pieces of code all result in the exact same computation graph when traced as shown:

+----------------------------------------+-----------------------------------------+-----------------------------------------+
|.. code-block:: python |.. code-block:: python |.. code-block:: python |
Expand All @@ -427,9 +427,9 @@ For example, the following 3 pieces of code all compile to the exact same comput
| | -1, 1, (3, 3)) | -1, 1, (3, 3)) |
| # input | b = ivy.zeros((3,)) | b = ivy.zeros((3,)) |
| x = ivy.array([1., 2., 3.]) | | |
| | # compile graph | # compile graph |
| # compile graph | graph = ivy.compile_graph( | graph = ivy.compile_graph( |
| net.compile_graph(x) | clean, x, w, b) | unclean, x, w, b) |
| | # trace graph | # trace graph |
| # trace graph | graph = ivy.trace_graph( | graph = ivy.trace_graph( |
| net.trace_graph(x) | clean, x, w, b) | unclean, x, w, b) |
| | | |
| # execute graph | # execute graph | # execute graph |
| net(x) | graph(x, w, b) | graph(x, w, b) |
Expand All @@ -439,8 +439,8 @@ For example, the following 3 pieces of code all compile to the exact same comput
:align: center
:width: 75%

This compilation is not restricted to just PyTorch.
Let's take another example, but compile to Tensorflow, NumPy, and JAX:
This tracing is not restricted to just PyTorch.
Let's take another example, but trace to Tensorflow, NumPy, and JAX:

+------------------------------------+
|.. code-block:: python |
Expand All @@ -454,7 +454,7 @@ Let's take another example, but compile to Tensorflow, NumPy, and JAX:
| x = ivy.array([[1., 2., 3.]]) |
| y = ivy.array([[2., 3., 4.]]) |
| # create graph |
| graph = ivy.compile_graph( |
| graph = ivy.trace_graph( |
| ivy_func, x, y) |
| |
| # call graph |
Expand Down Expand Up @@ -486,13 +486,13 @@ Jax:
:width: 75%
|
The example above further emphasizes that the graph compiler creates a computation graph consisting of backend functions, not Ivy functions.
Specifically, the same Ivy code compiles to different graphs depending on the selected backend.
However, when compiling native framework code, we are only able to compile a graph for that same framework.
For example, we cannot take torch code and compile this into tensorflow code.
The example above further emphasizes that the tracer creates a computation graph consisting of backend functions, not Ivy functions.
Specifically, the same Ivy code is traced to different graphs depending on the selected backend.
However, when compiling native framework code, we are only able to trace a graph for that same framework.
For example, we cannot take torch code and trace this into tensorflow code.
However, we can transpile torch code into tensorflow code (see `Ivy as a Transpiler <ivy_as_a_transpiler.rst>`_ for more details).

The graph compiler does not compile to C++, CUDA, or any other lower level language.
The tracer is not a compiler and does not compile to C++, CUDA, or any other lower level language.
It simply traces the backend functional methods in the graph, stores this graph, and then efficiently traverses this graph at execution time, all in Python.
Compiling to lower level languages (C++, CUDA, TorchScript etc.) is supported for most backend frameworks via :func:`ivy.compile`, which wraps backend-specific compilation code, for example:

Expand Down Expand Up @@ -524,6 +524,6 @@ Therefore, the backend code can always be run with maximal efficiency by compili

**Round Up**

Hopefully, this has painted a clear picture of the fundamental building blocks underpinning the Ivy framework, being the backend functional APIs, Ivy functional API, backend handler, and graph compiler 🙂
Hopefully, this has painted a clear picture of the fundamental building blocks underpinning the Ivy framework, being the Backend functional APIs, Ivy functional API, Backend handler, and Tracer 😄

Please reach out on `discord <https://discord.gg/sXyFF8tDtm>`_ if you have any questions!
2 changes: 1 addition & 1 deletion docs/overview/design/ivy_as_a_framework.rst
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
Ivy as a Framework
==================

On the `Building Blocks <building_blocks.rst>`_ page, we explored the role of the backend functional APIs, the Ivy functional API, the framework handler, and the graph compiler.
On the `Building Blocks <building_blocks.rst>`_ page, we explored the role of the Backend functional APIs, the Ivy functional API, the Backend handler, and the Tracer.
These are parts labeled as (a) in the image below.

On the `Ivy as a Transpiler <ivy_as_a_transpiler.rst>`_ page, we explained the role of the backend-specific frontends in Ivy, and how these enable automatic code conversions between different ML frameworks.
Expand Down
Loading

0 comments on commit 39671b3

Please sign in to comment.