docs: update formatting and examples (PR2472)

Author: kshitij12345

docs/source/fundamentals/installation.rst

@@ -56,11 +56,11 @@ Thunder can easily integrate OpenAI Triton kernels. You can install Triton using
 Install Thunder
 ===============
 
-You can now install Thunder
+You can now install Thunder::
 
   pip install git+https://github.com/Lightning-AI/lightning-thunder.git
 
-Alternatively you can clone the Thunder repository and install locally
+Alternatively you can clone the Thunder repository and install locally::
 
   git clone https://github.com/Lightning-AI/lightning-thunder.git
   cd lightning-thunder

docs/source/intermediate/additional_executors.rst

@@ -10,11 +10,9 @@ Triton CrossEntropy Executor
 
 The Triton CrossEntropy executor can execute ``torch.cross_entropy()`` using an optimized kernel written in OpenAI Triton (https://github.com/openai/triton). It can be used like in the following example::
 
+  import torch
   import thunder
-  from thunder.executors import nvfuserex, torchex
-  from thunder.executors.triton_crossentropy import deregister_triton_entropyex, register_triton_entropyex
-
-  register_triton_entropyex(add_to_default_executors=False)
+  from thunder.executors.triton_crossentropy import triton_ex as triton_cross_entropy_ex
 
   def xentropy(logits, labels, weight, reduction, ignore_index):
       return thunder.torch.cross_entropy(
@@ -23,7 +21,7 @@ The Triton CrossEntropy executor can execute ``torch.cross_entropy()`` using an
 
   jitted_xentropy = thunder.jit(
     xentropy,
-    executors_list=['triton_crossentropy', nvfuserex, torchex]
+    executors=[triton_cross_entropy_ex,]
   )
 
   device = 'cuda'
@@ -41,43 +39,42 @@ The Triton CrossEntropy executor can execute ``torch.cross_entropy()`` using an
 
 This prints::
 
-  # Constructed by Delete Last Used
+  # Constructed by Delete Last Used (took 0 milliseconds)
   import torch
+  from thunder.executors.torchex import no_autocast
+
   @torch.no_grad()
-  def xentropy(logits, labels, weight, reduction, ignore_index):
+  @no_autocast()
+  def computation(logits, labels, weight):
     # logits: "cuda:0 f32[2048, 50257]"
     # labels: "cuda:0 i64[2048]"
     # weight: "cuda:0 f32[50257]"
-    # "sum"
-    # ignore_index: "int 10106"
-    t22 = triton_cross_entropy(logits, labels, weight, None, ignore_index, None, "sum", 0.0)  # t22: "cuda:0 f32[]"
-    del [logits, labels, weight, ignore_index]
-    return t22
+    t23 = triton_crossentropy(logits, labels, weight, None, 45279, None, 'sum', 0.0)  # t23: "cuda:0 f32[]"
+    del logits, labels, weight
+    return t23
 
-As shown in the above trace, ``triton_cross_entropy()`` is the one running the operation.
+As shown in the above trace, ``triton_crossentropy()`` is the one running the operation.
 
 Apex CrossEntropy Executor
 ==========================
 
 The Apex CrossEntropy executor can execute ``torch.cross_entropy()`` through an optimized kernel, like this::
 
+  import torch
   import thunder
-  from thunder.executors import nvfuserex, torchex
-  from thunder.executors.apex_entropyex import deregister_apex_entropyex, register_apex_entropyex
-
-  register_apex_entropyex(add_to_default_executors=False)
+  from thunder.executors.apex_entropyex import apex_ex
 
   def xentropy(logits, labels):
       return thunder.torch.cross_entropy(
           logits, labels, reduction='mean', ignore_index=-1
       )
 
-  jitted_xentropy = thunder.jit(xentropy, executors_list=['apex_xentropy', nvfuserex, torchex])
+  jitted_xentropy = thunder.jit(xentropy, executors=[apex_ex,])
 
   device = 'cuda'
   dtype = torch.float32
 
-  logits = torch.randn([2048, 50257], device=device, dtype=thunder.torch.to_torch_dtype(dtype))
+  logits = torch.randn([2048, 50257], device=device, dtype=dtype)
   labels = torch.randint(0, 50257, [2048], device=device)
 
   jitted_xentropy(logits, labels)
@@ -86,14 +83,17 @@ The Apex CrossEntropy executor can execute ``torch.cross_entropy()`` through an
 
 This prints::
 
-  # Constructed by Delete Last Used
+  # Constructed by Delete Last Used (took 0 milliseconds)
   import torch
+  from thunder.executors.torchex import no_autocast
+
   @torch.no_grad()
-  def xentropy(logits, labels):
+  @no_autocast()
+  def computation(logits, labels):
     # logits: "cuda:0 f32[2048, 50257]"
     # labels: "cuda:0 i64[2048]"
-    t18 = apex_cross_entropy(logits, labels, None, None, -1, None, "mean", 0.0)  # t18: "cuda:0 f32[]"
-    del [logits, labels]
+    (t18, _) = apex_cross_entropy(logits, labels, 'mean', 0.0)
+    del logits, labels
     return t18
 
 showing that Apex is running the operation.