Formats code in documentation with black

This change updates our documentation generation to format all code blocks
with `black`. This helps ensure that lines do not overflow and create scrollable
elements in the rendered docs. For lines that `black` doesn't touch (e.g. comments),
we include our own assertions to ensure they don't exceed the length limits.

The logic to extract code from a code block (i.e. stripping out markup) has also
been consolidated.

Finally, this change removes the `manual` test cadence and related tests since
they were redundant (the same things are tested during doc generation) and didn't
add any value (the thought was they would make it easier to debug, but (1) they don't
and (2) it's rarely difficult to figure out what's wrong with the code blocks in documentation)

Author: pranavm-nvidia

tripy/docs/conf.py

@@ -286,10 +286,9 @@ def allow_no_example():
 
         code_block_lines, local_var_lines, output_lines, _ = helper.process_code_block_for_outputs_and_locals(
             block,
-            block.code(),
             format_contents=lambda title, contents, lang: f"\n\n.. code-block:: {lang}\n"
             + indent((f":caption: {title}" if title else "") + f"\n\n{contents}", prefix=" " * helper.TAB_SIZE),
-            err_msg=f"Failed while processing docstring for: {what}: {name} ({obj})",
+            err_msg=f"Failed while processing docstring for: {what}: {name} ({obj}): ",
             strip_assertions=True,
         )
 

tripy/docs/generate_rsts.py

@@ -208,9 +208,8 @@ def add_block(title, contents, lang):
             code_block_lines, local_var_lines, output_lines, code_locals = (
                 helper.process_code_block_for_outputs_and_locals(
                     block.raw_str(),
-                    str(block),
                     format_contents=add_block,
-                    err_msg=f"Error while executing code block from {guide_path}.",
+                    err_msg=f"Error while executing code block {index} (line {block.line_number}) from {guide_path}. ",
                     local_vars=code_locals,
                 )
             )

tripy/docs/post0_developer_guides/how-to-add-new-ops.md

@@ -47,9 +47,10 @@ from tripy.flat_ir.ops.base import BaseFlatIROp
 class ThetaOp(BaseFlatIROp):
     dim: int
 
-    # `to_mlir()` is the trickiest bit. As the name implies, the method is meant to lower the
-    # `FlatIR` operator into MLIR. To figure out which MLIR operators to use, refer to
-    # the 'MLIR Python API Guide' (linked below).
+    # `to_mlir()` is the trickiest bit. As the name implies, the method is
+    # meant to lower the `FlatIR` operator into MLIR. To figure out which
+    # MLIR operators to use, refer to the 'MLIR Python API Guide'
+    # (linked below).
     def to_mlir(self, operands):
         out_type = self.outputs[0].to_mlir()
         theta_dim = ir.IntegerAttr.get(type=ir.IntegerType.get_signless(64), value=self.dim)
@@ -116,29 +117,31 @@ from tripy.frontend.trace.ops.base import BaseTraceOp
 import tripy.frontend.trace.ops.utils as op_utils
 
 
-# Just like with `FlatIR` operators, all `Trace` operators are implemented as `dataclass`es.
-# As before, we want `repr=False` here.
+# Just like with `FlatIR` operators, all `Trace` operators are implemented
+# as `dataclass`es. As before, we want `repr=False` here.
 @dataclass(repr=False)
 class Theta(BaseTraceOp):
-    # Notice that we do *not* need to define a constructor and can rely on the default
-    # implementation provided by `dataclass`.
+    # Notice that we do *not* need to define a constructor and can rely on
+    # the default implementation provided by `dataclass`.
     dim: int
     dtype: datatype.dtype
 
     # `infer_rank()` populates the rank of the output `TraceTensor`s.
-    # Here we use one of the predefined policies to set the output rank to the same as the shape (i.e. the length)
-    # of the shape operand.
+    # Here we use one of the predefined policies to set the output rank
+    # to the same as the shape (i.e. the length) of the shape operand.
     infer_rank = op_utils.InferRankPolicies.same_as_shape_of_shape_input()
 
     # *Optional* `infer_dtypes()` populates the data types of the
     # output `TraceTensor`s. The default implementation copies the input
-    # data types if they are all the same, so you may not need to implement this.
+    # data types if they are all the same, so you may not need to implement
+    # this.
     def infer_dtypes(self):
         self.outputs[0].dtype = self.dtype
 
     # *Optional* `infer_devices()` populates the devices of the
     # output `TraceTensor`s. The default implementation copies the input
-    # devices if they are all the same, so you may not need to implement this either.
+    # devices if they are all the same, so you may not need to implement
+    # this either.
     def infer_devices(self):
         self.outputs[0].device = device("gpu")
 
@@ -177,30 +180,35 @@ from tripy import export
 import tripy.frontend.utils as frontend_utils
 from tripy.types import ShapeLike
 
-# We can use the `export.public_api()` decorator to automatically export this function into the
-# top-level module. This means it will be accessible as `tripy.theta`.
+# We can use the `export.public_api()` decorator to automatically export this
+# function into the top-level module. This means it will be accessible as
+# `tripy.theta`.
 #
-# This decorator also controls how the API is exposed in the documentation - the `document_under`
-# option determines where in the documentation hierarchy this API will show up.
+# This decorator also controls how the API is exposed in the documentation -
+# the `document_under` option determines where in the documentation hierarchy
+# this API will show up.
 #
-# If we needed to provide any special autodoc options, we could use the `autodoc_options` parameter.
+# If we needed to provide any special autodoc options, we could use the
+# `autodoc_options` parameter.
 @export.public_api(document_under="tensor_operations")
 
-# The `convert_to_tensors` decorator automatically converts compatible arguments,
-# like `TensorLike` or `ShapeLike`s, into tensors.
+# The `convert_to_tensors` decorator automatically converts compatible
+# arguments, like `TensorLike` or `ShapeLike`s, into tensors.
 @frontend_utils.convert_to_tensors()
 def theta(shape: ShapeLike, dim: int = 0, dtype: datatype.dtype = datatype.float32) -> "tripy.Tensor":
-    # For any public facing interfaces, we have documentation requirements which you can read
-    # about in the 'Docs README' (linked below). The docstring we've implemented here
-    # adheres to all of these requirements. Non-compliant docstrings will, in most cases,
-    # cause test failures; however, you should still manually ensure you're writing high-quality
-    # docstrings.
+    # For any public facing interfaces, we have documentation requirements which
+    # you can read about in the 'Docs README' (linked below). The docstring
+    # we've implemented here adheres to all of these requirements. Non-compliant
+    # docstrings will, in most cases, cause test failures; however, you should
+    # still manually ensure you're writing high-quality docstrings.
     #
-    # The examples in docstrings are run as part of our tests, so you should also add
-    # assertions to make sure things are functionally correct. In this case, we check
-    # that the `output` we create in the code example is what we expect.
+    # The examples in docstrings are run as part of our tests, so you should
+    # also add assertions to make sure things are functionally correct. In this
+    # case, we check that the `output` we create in the code example is what we
+    # expect.
     """
-    Fills an output tensor with consecutive values starting from zero along the given dimension.
+    Fills an output tensor with consecutive values starting from zero
+    along the given dimension.
 
     Args:
         shape: The desired shape.
@@ -217,12 +225,15 @@ def theta(shape: ShapeLike, dim: int = 0, dtype: datatype.dtype = datatype.float
 
         output = tp.theta([3])
 
-        assert np.array_equal(cp.from_dlpack(output).get(), np.arange(0, 3, dtype=np.float32))
+        assert np.array_equal(
+            cp.from_dlpack(output).get(), np.arange(0, 3, dtype=np.float32)
+        )
     """
 
-    # Next we build the trace operator. The `build()` function is also responsible for constructing
-    # the output frontend Tensors. All of the arguments that follow the inputs
-    # are forwarded directly to the constructor of the `Trace` operator.
+    # Next we build the trace operator. The `build()` function is also
+    # responsible for constructing the output frontend Tensors. All of the
+    # arguments that follow the inputs are forwarded directly to the
+    # constructor of the `Trace` operator.
     return Theta.build([shape], dim, dtype)
 
 ```

tripy/pyproject.toml

@@ -40,6 +40,7 @@ build = [
   "mypy==1.11.0",
 ]
 doc_test_common = [
+  "black==24.10.0",
   "torch==2.4.0+cu121",
   "numpy==1.25.0",
   # cupy requires NVRTC but does not specify it as a package dependency
@@ -96,5 +97,4 @@ testpaths = [
 addopts = "--strict-markers"
 markers = [
     "l1: Indicates that the test should only be run in nightlies.",
-    "manual: Disables tests in automation",
 ]

tripy/tests/README.md

@@ -18,13 +18,12 @@ You can also provide marker arguments to only run specific test cadences
 L0 tests, use:
 
 ```bash
-pytest tests/ -v -m "not l1 and not manual" -n 4 --dist worksteal --ignore tests/performance
-pytest tests/performance -v -m "not l1 and not manual"
+pytest tests/ -v -m "not l1" -n 4 --dist worksteal --ignore tests/performance
+pytest tests/performance -v -m "not l1"
 ```
 
-Note that the L0/L1 tests can be parallelized, which is not necessarily
-true of `manual` tests. In that case, performance tests are run separately
-because they must run serially to ensure accurate measurements.
+Note that the L0/L1 tests can be parallelized. In that case, performance tests
+are run separately because they must run serially to ensure accurate measurements.
 
 ## Profiling
 
@@ -36,7 +35,7 @@ tests together.
 For example, to profile L0 tests, run:
 
 ```bash
-pytest tests/ -v -m "not l1 and not manual" --ignore tests/performance --profile
+pytest tests/ -v -m "not l1" --ignore tests/performance --profile
 ```
 
 You can visualize the results using `snakeviz`.