v4.4 update. (#2979)

This commit is contained in:
Junkai-Wu
2026-01-25 00:46:17 +08:00
committed by GitHub
parent 2fafefb7b9
commit 9fba3195f9
293 changed files with 46344 additions and 2996 deletions

View File

@@ -17,4 +17,6 @@ CuTe DSL
Debugging with the DSL <cute_dsl_general/debugging.rst>
Autotuning with the DSL <cute_dsl_general/autotuning_gemm.rst>
Educational Notebooks <cute_dsl_general/notebooks.rst>
Deprecation Policy <deprecation.rst>
Compile with TVM FFI <cute_dsl_general/compile_with_tvm_ffi.rst>
Ahead-of-Time (AOT) Compilation <cute_dsl_general/dsl_ahead_of_time_compilation.rst>

View File

@@ -74,7 +74,7 @@ Changelog for CuTe DSL API changes
- Introduce BlockScaled layout utilities in `blockscaled_layout.py <https://github.com/NVIDIA/cutlass/blob/main/python/CuTeDSL/cutlass/utils/blockscaled_layout.py>`_ for creating the required scale factor layouts in global memory, shared memory and tensor memory.
* ``cutlass.cute.compile`` now supports compilation options. Refer to `JIT compilation options <https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/cute_dsl_general/dsl_jit_compilation_options.html>`_ for more details.
* ``cutlass.cute.testing.assert_`` now works for device JIT function. Specify ``--enable-device-assertions`` as compilation option to enable.
* ``cutlass.cute.testing.assert_`` now works for device JIT function. Specify ``--enable-assertions`` as compilation option to enable.
* ``cutlass.cute.make_tiled_copy`` is now deprecated. Please use ``cutlass.cute.make_tiled_copy_tv`` instead.
* Shared memory capacity query

View File

@@ -4,7 +4,8 @@
Compile with TVM FFI
====================
Apache TVM FFI is an open ABI and FFI for machine learning systems. More information can be found in the `official documentation <https://tvm.apache.org/ffi/>`_.
Apache TVM FFI is an open ABI and FFI for machine learning systems. More information can be found in
the `official documentation <https://tvm.apache.org/ffi/>`_.
To install TVM FFI, you can run the following command:
@@ -14,7 +15,9 @@ To install TVM FFI, you can run the following command:
# optional package for improved torch tensor calling performance
pip install torch-c-dlpack-ext
In |DSL|, TVM FFI can be enabled as an option for JIT-compiled functions. Using TVM FFI can lead to faster JIT function invocation and provides better interoperability with machine learning frameworks (e.g., directly take ``torch.Tensor`` as arguments).
In |DSL|, TVM FFI can be enabled as an option for JIT-compiled functions. Using TVM FFI can lead to faster
JIT function invocation and provides better interoperability with machine learning frameworks
(e.g., directly take ``torch.Tensor`` as arguments).
Enable Apache TVM FFI in |DSL|
@@ -129,7 +132,8 @@ stride via the ``stride`` argument in the ``make_fake_tensor`` API.
``cute.Tensor`` adapter for TVM FFI
-----------------------------------
To adapt the ``cute.Tensor`` to the TVM FFI function, you can use the ``cute.runtime.from_dlpack`` function with the ``enable_tvm_ffi=True`` option or the environment variable ``CUTE_DSL_ENABLE_TVM_FFI=1``. For example:
To adapt the ``cute.Tensor`` to the TVM FFI function, you can use the ``cute.runtime.from_dlpack`` function with the
``enable_tvm_ffi=True`` option or the environment variable ``CUTE_DSL_ENABLE_TVM_FFI=1``. For example:
.. code-block:: python
@@ -570,7 +574,7 @@ Then you can load back the exported module and use it in different ways:
from cutlass import cute
def example_load_module_add_one():
mod = cute.runtime.load_module("./add_one.so")
mod = cute.runtime.load_module("./add_one.so", enable_tvm_ffi=True)
a_torch = torch.arange(10, dtype=torch.float32, device="cuda")
b_torch = torch.empty(10, dtype=torch.float32, device="cuda")
mod.add_one(a_torch, b_torch)
@@ -587,9 +591,11 @@ in the official documentation.
When you build your own libraries, make sure you link against the necessary runtime libraries.
You can use ``cute.runtime.find_runtime_libraries(enable_tvm_ffi=True)`` to get the path to these libraries.
``cute.runtime.load_module`` will load these libraries automatically before loading
``cute.runtime.load_module(path, enable_tvm_ffi=True)`` will load these libraries automatically before loading
an exported module. You can also manually load these libraries in advanced use cases.
For low-level cute ABI AOT compilation support without TVM FFI, you can refer to :doc:`dsl_ahead_of_time_compilation`.
Keyword Arguments and Defaults
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -683,3 +689,32 @@ The code block below shows how to do this:
wrapped_func(a_torch, b_torch, offset=4)
print("result of b_torch after wrapped_func(a_torch, b_torch, offset=4)")
print(b_torch)
Limitations
-----------
The Fake Tensor flow is ONLY compatible with TVM FFI because TVM FFI supports more flexible constraints on Tensor arguments.
For instance, fake tensor can specify per-mode static shape or constraints on shape and strides which are not supported by
existing ``from_dlpack`` flow. It's expected that JIT function compiled with fake tensor will have different ABI compared to
tensor converted by ``from_dlpack``.
.. code-block:: python
import cutlass.cute as cute
import torch
n = cute.sym_int()
# Dynamic Shape
fake_a = cute.runtime.make_fake_compact_tensor(cute.Float32, (n,))
# Compile without tvm-ffi
compiled_fn = cute.compile(foo, fake_a)
# Wrong, in compatible ABI
compiled_fn(from_dlpack(a))
In order to avoid such issue, it's recommended to use fake tensor only with TVM FFI backend. Practically speaking,
as we only want to call ``from_dlpack`` once and reuse for both compilation and runtime, the benefit of
using fake tensor is limited in this case.

View File

@@ -172,6 +172,25 @@ For detecting memory errors and race conditions:
Please refer to the `compute-sanitizer documentation <https://developer.nvidia.com/compute-sanitizer>`_ for more details.
Set function name prefix
~~~~~~~~~~~~~~~~~~~~~~~~~
By default, the function name (host function or kernel function) is automatically generated based on the function name and its parameters.
Sometimes you may want to attach some runtime information to the function name to make performance profiling and debugging easier,
e.g., the kernel configs or the rank ids. You can assign a name prefix to the name by calling the ``set_name_prefix``
method on the host function or kernel function.
.. code:: python
@cute.kernel
def kernel(arg1, arg2, ...):
...
@cute.jit
def launch_kernel():
kernel.set_name_prefix("your_custom_name_prefix")
kernel(arg1, arg2, ...).launch(grid=[1, 1, 1], block=[1, 1, 1], ...)
For above example, the generated kernel name will be "your_custom_name_prefix_xxx".
Conclusion
----------

View File

@@ -0,0 +1,239 @@
.. _dsl_ahead_of_time_compilation:
.. |DSL| replace:: CuTe DSL
Ahead-of-Time (AOT) Compilation
===============================
This guide demonstrates how to use |DSL|'s Ahead-of-Time (AOT) compilation features to export compiled kernels for use in production environments.
Overview
--------
|DSL| Ahead-of-Time (hereinafter referred to as AOT) compilation allows you to:
* **Compile once, enable cross-compilation**: Write kernels in Python and cross-compile them for multiple GPU architectures.
* **Remove JIT overhead**: Eliminate compilation delays in production by pre-compiling kernels.
* **Flexible integration**: Easily integrate compiled kernels into both Python and C/C++ codebases using flexible deployment options.
We provide 2 levels of AOT ABI:
1. **Low-Level CuTe ABI**: This ABI is expressed using CuTe DSL types and tensors, mirroring the original Python function.
2. **High-Level Apache TVM FFI ABI**: For interop with various frameworks (e.g., PyTorch, JAX), and offer high-level stable ABI access.
This guide will focus on the CuTe ABI AOT. For the Apache TVM FFI AOT, please refer to the section "Exporting Compiled Module" in :doc:`compile_with_tvm_ffi`.
CuTe ABI AOT Workflow
---------------------
Export Interface
~~~~~~~~~~~~~~~~
The ``export_to_c`` interface is provided by the ``JitCompiledFunction`` class. It accepts the following parameters:
* ``file_path``: The path to the directory where the header and object files will be saved.
* ``file_name``: The base name for the header and object files. The same file name will always overwrite existing files.
* ``function_prefix``: The prefix of the function symbol in the generated object file. This should be a unique identifier to avoid symbol conflicts. Users should ensure the function prefix is unique for each exported function. Defaults to the ``file_name``.
It generates the following files:
* ``{file_path}/{file_name}.h``: A C header file containing API function declarations. This header specifies the runtime function signatures in C, mirroring the original Python function interfaces.
* ``{file_path}/{file_name}.o``: A standard object file containing the compiled kernel code. You can link this object file into either a static or shared library. It includes the host entry function, fatbin data, and helper functions such as ``cuda_init`` and ``cuda_load_to_device``. Additionally, it embeds metadata for runtime loading and version verification.
Example:
.. code-block:: python
import cutlass.cute as cute
import cutlass.cute.cuda as cuda
@cute.kernel
def print_tensor_kernel(a: cute.Tensor):
cute.printf("a: {}", a)
@cute.jit
def print_tensor(a: cute.Tensor, stream: cuda.CUstream):
print_tensor_kernel(a).launch(grid=(1, 1, 1), block=(1, 1, 1), stream=stream)
compiled_func = cute.compile(print_tensor)
# Export compiled functions to object files and headers
compiled_func.export_to_c(file_path="./artifacts", file_name="print_tensor_example", function_prefix="print_tensor")
Loading in Python
~~~~~~~~~~~~~~~~~
Load pre-compiled object files or shared libraries into Python for execution.
.. code-block:: python
import cutlass.cute as cute
import torch
from cutlass.cute import from_dlpack
import cutlass.cute.cuda as cuda
# Load module from object file
module = cute.runtime.load_module("./artifacts/print_tensor_example.o")
# or
module = cute.runtime.load_module("./artifacts/libprint_tensor_example.so")
# Prepare data
a = torch.arange(160, dtype=torch.float32, device="cuda").reshape(16, 10)
a_cute = from_dlpack(a).mark_layout_dynamic()
stream = cuda.CUstream(0)
# Call the function (no JIT compilation needed!)
module.print_tensor(a_cute, stream=stream)
# This will fail because 'non_existing_api' was not exported:
# module.non_existing_api()
C++ Integration with Static Linking
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Integrate compiled kernels directly into your C++ executable during the build process. The generated header file supplies the necessary API for loading the module and invoking the function.
Example:
.. code-block:: cpp
#include "print_tensor_example.h"
#include <cuda_runtime.h>
void run_print_tensor() {
// Prepare tensor, the tensor declaration is in the header file
print_tensor_Tensor_a_t tensor_a;
tensor_a.data = nullptr; // GPU memory is set to nullptr.
// Set dynamic shapes and strides
tensor_a.dynamic_shapes[0] = 32;
tensor_a.dynamic_shapes[1] = 16;
tensor_a.dynamic_strides[0] = 16;
// Create stream
cudaStream_t stream;
cudaStreamCreate(&stream);
// Load module before calling the kernel
print_tensor_Kernel_Module_t module;
print_tensor_Kernel_Module_Load(&module);
// Call the kernel; the kernel wrapper function is defined in the header file
cute_dsl_print_tensor_wrapper(&module, &tensor_a, stream);
// Cleanup
print_tensor_Kernel_Module_Unload(&module);
cudaStreamDestroy(stream);
}
The ``print_tensor_example.h`` header file is generated by the ``export_to_c`` interface. It includes:
* The ``print_tensor_Kernel_Module_t`` type: Represents the kernel module.
* The ``print_tensor_Tensor_a_t`` type: A tensor-specific type that defines the ABI for a particular CuTe tensor.
* The ``cute_dsl_print_tensor_wrapper`` function: The user-facing entry point to invoke the kernel.
The compilation of the C++ executable requires the ``libcuda_dialect_runtime.so`` or ``libcuda_dialect_runtime_static.a`` library which is involved in ``<wheel_install_path>/lib``, along with the CUDA driver and runtime libraries, to function properly.
C++ Integration with Dynamic Loading
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Dynamically load pre-compiled object files or shared libraries at runtime. By including the ``CuteDSLRuntime.h`` header, you can load the module, look up exported functions, and invoke them.
.. code-block:: cpp
#include "CuteDSLRuntime.h"
#include <cuda_runtime.h>
void run_print_tensor() {
// Load module from shared library
CuteDSLRT_Module_t *module = nullptr;
CuteDSLRT_Error_t err = CuteDSLRT_Module_Load(
&module,
"./artifacts/libprint_tensor_example.so"
);
// or
CuteDSLRT_Error_t err = CuteDSLRT_Module_Load(
&module,
"./artifacts/print_tensor_example.o"
);
check_error(err);
// Lookup function
CuteDSLRT_Function_t *func = nullptr;
err = CuteDSLRT_Module_Get_Function(&func, module, "print_tensor");
check_error(err);
// Prepare arguments, matching the argument type defined in the header file
typedef struct {
void *data;
int32_t dynamic_shapes[2];
int64_t dynamic_strides[1];
} print_tensor_Tensor_a_t;
print_tensor_Tensor_a_t tensor_a;
tensor_a.data = nullptr;
tensor_a.dynamic_shapes[0] = 32;
tensor_a.dynamic_shapes[1] = 16;
tensor_a.dynamic_strides[0] = 16;
// Create stream
cudaStream_t stream;
cudaStreamCreate(&stream);
// Call the function; the runtime function accepts packed arguments, refer to the wrapper in the header file
int ret;
void* args[] = {&tensor_a, &stream, &ret};
err = CuteDSLRT_Function_Run(func, args, 3);
check_error(err);
cudaStreamSynchronize(stream);
// Cleanup
CuteDSLRT_Module_Destroy(module);
cudaStreamDestroy(stream);
}
The ``CuteDSLRuntime.h`` header file can be found in ``<wheel_install_path>/include``. It includes:
* The ``CuteDSLRT_Error_t`` type: Indicates error status.
* The ``CuteDSLRT_Module_Load`` function: Loads the module.
* The ``CuteDSLRT_Module_Get_Function`` function: Gets a function from the loaded module. The runtime API will load the CUDA module for kernel execution.
* The ``CuteDSLRT_Function_Run`` function: Runs the function.
* The ``CuteDSLRT_Module_Destroy`` function: Destroys the module.
The compilation of the C++ executable requires the ``libcute_dsl_runtime.so`` library which is involved in ``<wheel_install_path>/lib``, along with the CUDA driver and runtime libraries, to function properly.
Supported Argument Types
------------------------
|DSL| supports the following argument types:
* ``cute.Tensor``
* ``cute.Shape`` / ``cute.Coord`` / ``cute.Tile`` / ``cute.IntTuple`` / ``cute.Stride``
* ``cuda.CUstream``
* ``cutlass.Int8`` / ``cutlass.Int16`` / ``cutlass.Int32`` / ``cutlass.Int64`` / ``cutlass.Boolean``
* ``cutlass.Uint8`` / ``cutlass.Uint16`` / ``cutlass.Uint32`` / ``cutlass.Uint64``
* ``cutlass.Float32`` / ``cutlass.TFloat32`` / ``cutlass.Float64`` / ``cutlass.Float16``
Note that:
1. ``cute.Tensor`` is a dynamic tensor type that only contains dynamic shapes and strides in its ABI representation. As a result, different compilations may produce different tensor ABIs. This is why declarations for each tensor type are included in the generated header file.
2. ``strides`` in ``cute.Tensor`` are determined by the ``use_32bit_strides`` compile argument. When ``use_32bit_strides`` is set to ``True``, the strides are 32-bit; when set to ``False``, they are 64-bit.
3. Currently, custom types are not supported for AOT compilation.
Object File Compatibility Issues
--------------------------------
The object file generated by |DSL| depends on the CUDA runtime library. Therefore, ensure that the version of the CUDA runtime/toolkit library matches the version used by |DSL|. Otherwise, ABI compatibility with the CUDA runtime cannot be guaranteed.
When using C++ static linking integration, compatibility is assured because the header and object files are generated together and guaranteed to match.
For C++ dynamic loading integration and Python loading, the binary file is loaded at runtime. To ensure compatibility, version information is embedded in the metadata of the generated binary file. At runtime, this version information is checked, and if it does not match the expected version, the binary file will be rejected.
Relation to Apache TVM FFI AOT
------------------------------
Apache TVM FFI AOT offers a comparable capability, enabling TVM functions to be compiled into binary files that can be loaded and executed at runtime.
For more information, see the section "Exporting Compiled Module" in :doc:`compile_with_tvm_ffi`.
The primary distinction is that, when TVM FFI is enabled, |DSL| generates a dedicated wrapper function on top of the underlying CuTe ABI. This wrapper adheres to the calling conventions defined by TVM FFI.
In contrast, the CuTe ABI entry function is specified directly in the generated header file, which affects how arguments must be provided.
For instance, with the TVM FFI wrapper function, users are able to pass in arguments such as ``torch.Tensor`` directly. However, when calling the CuTe ABI entry function, arguments should be provided as ``cute.Tensor`` types.

View File

@@ -184,6 +184,46 @@ Standard Python ``while`` is supported.
n += 1
Summary of Control Flow behavior
---------------------------------
.. list-table::
:header-rows: 1
:widths: 30 30 30
* - **Control Flow**
- **Run time evaluation**
- **Compile time evaluation**
* - if cutlass.const_expr()
-
-
* - if pred
-
-
* - while cutlass.const_expr()
-
-
* - while pred
-
-
* - for i in cutlass.range_constexpr()
-
-
* - for i in range()
-
-
* - for i in cutlass.range() (support advanced unrolling and pipelining)
-
-
Compile-Time Metaprogramming
----------------------------

View File

@@ -73,7 +73,7 @@ You can use the following code to specify compilation options:
jit_executor_with_opt_level_2 = cute.compile(add, 1, 2, options="--opt-level 2")
jit_executor_with_opt_level_1 = cute.compile(add, 1, 2, options="--opt-level 1")
jit_executor_with_enable_device_assertions = cute.compile(add, 1, 2, options="--enable-assertions")
jit_executor_with_enable_assertions = cute.compile(add, 1, 2, options="--enable-assertions")
jit_executor_with_keep_cubin = cute.compile(add, 1, 2, options="--keep-cubin")
jit_executor_with_keep_ptx = cute.compile(add, 1, 2, options="--keep-ptx")
jit_executor_with_ptxas_options = cute.compile(add, 1, 2, options="--ptxas-options '--opt-level=2'")
@@ -100,7 +100,7 @@ Notebly, boolean options are automatically converted to True instances of the op
jit_executor_with_opt_level_2 = cute.compile[OptLevel(2)](add, 1, 2)
jit_executor_with_opt_level_1 = cute.compile[OptLevel(1)](add, 1, 2)
jit_executor_with_enable_device_assertions = cute.compile[EnableAssertions](add, 1, 2)
jit_executor_with_enable_assertions = cute.compile[EnableAssertions](add, 1, 2)
jit_executor_with_keep_cubin = cute.compile[KeepCUBIN](add, 1, 2)
jit_executor_with_keep_ptx = cute.compile[KeepPTX](add, 1, 2)
jit_executor_with_ptxas_options = cute.compile[PtxasOptions("--opt-level=2")](add, 1, 2)
jit_executor_with_ptxas_options = cute.compile[PtxasOptions("--opt-level=2")](add, 1, 2)

View File

@@ -0,0 +1,56 @@
.. _deprecation:
Deprecation Policy
==================
Purpose
-------
The goal of this policy is to evolve the DSL and its APIs while keeping user
programs stable. Features or APIs are deprecated only when they are redundant,
unsafe, or block better designs.
Deprecation Process
-------------------
**Step 1 — Soft Deprecation**
When a feature is considered for removal, it is first annotated with the
``@deprecated`` decorator or ``DeprecationWarning`` and documented with a
suggested alternative. At this stage, the feature continues to work normally.
Users are encouraged to provide feedback and describe their use cases.
If there is strong justification, we may keep or redesign the feature.
**Step 2 — Removal (the subsequent release)**
If no valid use cases remain, the deprecated feature will be removed in the
following **minor** release.
.. note::
The release version follows the format ``<major>.<minor>.<patch>``.
Communication
-------------
All deprecations are announced through:
* This page
* In-code warning messages
Soft Deprecations
-----------------
**Version 4.2.1**
* ``cute.arch.warpgroup_reg_alloc`` and ``cute.arch.warpgroup_reg_dealloc``
→ Scheduled for deprecation. Use ``cute.arch.setmaxregister_increase`` and ``cute.arch.setmaxregister_decrease`` instead.
* ``alignment`` argument in ``CooperativeGroup`` constructor
→ Scheduled for deprecation. It was unused; no replacement is suggested.
Deprecated Features
-------------------
*(None currently.)*

View File

@@ -17,12 +17,8 @@ the DSL.
Notable unsupported features
----------------------------
- Programmatic Dependent Launch (PDL)
- convolutions
- full support for ahead of time compilation
- preferred clusters
- CLC-based tile schedulers
- EVT support
- Windows support
Programming Model
@@ -82,16 +78,30 @@ Programming Model
xs.append(Float32(1.0))
**Python Function**
The DSL currently does not implement support for return values from Python functions,
although this capability is planned for future releases.
The DSL currently has **limited support for return values** from Python functions.
At the moment, only ``constexpr`` values can be returned, while returning **dynamic values** is **not yet supported**.
This capability is planned for a future release.
Example:
.. code:: python
.. code-block:: python
@cute.jit
def foo():
return 1 # Currently unsupported in CuTe DSL
def baz(a: cutlass.Constexpr):
return a + 1
@cute.jit
def foo(a: cutlass.Int32):
return a + 1
@cute.jit
def bar(a: cutlass.Int32):
val = foo(a) # works
val = baz(10) # works
val = bar(10) # works
foo(10) # currently unsupported in CuTe DSL
**Expression or Statement with Dependent Type**
CuTe DSL implements static typing and does not support dependent types.

View File

@@ -105,4 +105,4 @@ You can:
- Propose support for additional data types or kernel variants
- Help prioritize roadmap features by upvoting GitHub issues
Thank you for helping shape the future of CUTLASS DSLs!
Thank you for helping shape the future of CUTLASS DSLs!

View File

@@ -3,28 +3,39 @@
Quick Start Guide
=======================
The CUTLASS DSL 4.0 release currently supports **Linux** and **Python 3.12** only. To install CUTLASS DSLs (limited to CuTe DSL for now), use the following command
The CUTLASS DSL 4.4 release currently supports **Linux** and **Python 3.10 - 3.13** only. To install CUTLASS DSLs (limited to CuTe DSL for now), use the following command
Installation
-----------------------
To ensure compatibility with the examples and code on `GitHub <https://github.com/NVIDIA/cutlass/tree/main>`_,
use the `requirements.txt <https://github.com/NVIDIA/cutlass/blob/main/python/CuTeDSL/requirements.txt>`_ file from the corresponding commit in the repository.
use the `setup.sh <https://github.com/NVIDIA/cutlass/blob/main/python/CuTeDSL/setup.sh>`_ file from the corresponding commit in the repository.
.. code-block:: bash
git clone https://github.com/NVIDIA/cutlass.git
pip install -r cutlass/python/CuTeDSL/requirements.txt
If you just want to try out the last known stable release of the CUTLASS DSL (may not compatible with the latest examples and code), run:
git clone https://github.com/NVIDIA/cutlass.git
# For CUDA Toolkit 12.9:
./cutlass/python/CuTeDSL/setup.sh --cu12
# For CUDA Toolkit 13.1:
./cutlass/python/CuTeDSL/setup.sh --cu13
If you just want to try out the last known stable release of the CUTLASS DSL (may not be compatible with the latest examples and code), run:
.. code-block:: bash
# For CUDA Toolkit 12.9:
pip install nvidia-cutlass-dsl
# For CUDA Toolkit 13.1:
pip install nvidia-cutlass-dsl[cu13]
The ``nvidia-cutlass-dsl`` wheel includes everything needed to generate GPU kernels. It requires
the same NVIDIA driver version as the
`CUDA Toolkit 12.9 <https://docs.nvidia.com/cuda/cuda-toolkit-release-notes/index.html>`_.
the same NVIDIA driver version as the corresponding `CUDA Toolkit <https://docs.nvidia.com/cuda/cuda-toolkit-release-notes/index.html>`_
(CUDA Toolkit 12.9 or CUDA Toolkit 13.1).
Recommended Dependencies
---------------------------------
@@ -35,6 +46,8 @@ To run examples and begin development, we recommend installing:
pip install torch jupyter
We recommend installing JAX with CUDA support at version 0.8.1 to run JAX examples.
Recommended Python environment variables for jupyter notebooks
--------------------------------------------------------------