Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 24 additions & 0 deletions tripy/docs/conf.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
import sys

ROOT_DIR = os.path.join(os.path.abspath(os.path.dirname(__file__)), os.path.pardir)
import sphinx_toolbox
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could we group this with the other imports?

sys.path.insert(0, ROOT_DIR)
import inspect
import re
Expand Down Expand Up @@ -290,12 +291,35 @@ def process_signature(app, what, name, obj, options, signature, return_annotatio
return None


def visit_collapse_node(self, node):
label = getattr(node, 'label', None)

if label:
label = label.strip()
if label:
self.add_text(f"\n{label}:\n")

self.new_state(indent=0)


def depart_collapse_node(self, node):
self.end_state()


def setup(app):
# A note on aliases: if you rename a class via an import statement, e.g. `import X as Y`,
# the documentation generated for `Y` will just be: "Alias of X"
# To get the real documentation, you can make Sphinx think that `Y` is not an alias but instead a real
# class/function. To do so, you just need to define the __name__ attribute in this function (*not* in tripy code!):
# Y.__name__ = "Y"

# This is used to ensure that sphinx text build does not fail when it encounters a collapse node.
app.add_node(
sphinx_toolbox.collapse.CollapseNode,
html=(sphinx_toolbox.collapse.visit_collapse_node, sphinx_toolbox.collapse.depart_collapse_node),
override=True,
text=(visit_collapse_node, depart_collapse_node),
)

app.connect("autodoc-process-docstring", process_docstring)
app.connect("autodoc-process-signature", process_signature)
4 changes: 2 additions & 2 deletions tripy/nvtripy/frontend/ops/binary/pow.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@
@register_tensor_method("__pow__")
@wrappers.interface(
dtype_constraints={"self": "T1", "other": "T1", wrappers.RETURN_VALUE: "T1"},
dtype_variables={"T1": ["float32", "float16", "bfloat16", "int8"]},
dtype_variables={"T1": ["float32", "float16", "bfloat16", "int8", "int64"]},
convert_to_tensors=True,
)
def __pow__(self: "nvtripy.Tensor", other: TensorLike) -> "nvtripy.Tensor":
Expand Down Expand Up @@ -52,7 +52,7 @@ def __pow__(self: "nvtripy.Tensor", other: TensorLike) -> "nvtripy.Tensor":
@register_tensor_method("__rpow__")
@wrappers.interface(
dtype_constraints={"self": "T1", "other": "T1", wrappers.RETURN_VALUE: "T1"},
dtype_variables={"T1": ["float32", "float16", "bfloat16", "int8"]},
dtype_variables={"T1": ["float32", "float16", "bfloat16", "int8", "int64"]},
convert_to_tensors=True,
)
def __rpow__(self: "nvtripy.Tensor", other: TensorLike) -> "nvtripy.Tensor":
Expand Down
54 changes: 27 additions & 27 deletions tripy/pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ dependencies = [
"tensorrt>=10.11",
"mlir-tensorrt-compiler==0.1.42+cuda12.trt109",
"mlir-tensorrt-runtime==0.1.42+cuda12.trt109",
"colored==2.2.3",
"colored==2.3.0",
]

[project.urls]
Expand All @@ -21,20 +21,20 @@ Documentation = "https://nvidia.github.io/TensorRT-Incubator/"

[build-system]
requires = [
"setuptools==75.3.0",
"wheel==0.44.0",
"setuptools==80.9.0",
"wheel==0.45.1",
# For stubgen:
"mypy==1.11.0",
"mypy==1.16.1",
]
build-backend = "setuptools.build_meta"

[project.optional-dependencies]
dev = ["pre-commit==3.6.0"]
dev = ["pre-commit==4.2.0"]
build = [
"setuptools==75.3.0",
"wheel==0.44.0",
"setuptools==80.9.0",
"wheel==0.45.1",
# For stubgen:
"mypy==1.11.0",
"mypy==1.16.1",
]
doc_test_common = [
"black==24.10.0",
Expand All @@ -45,38 +45,38 @@ doc_test_common = [
"cupy-cuda12x",
]
docs = [
"sphinx==7.2.6",
"sphinx==7.4.7",
"furo==2024.8.6",
"sphinx-copybutton==0.5.2",
"sphinx-toolbox==3.5.0",
"docutils==0.20.1",
"myst-parser==2.0.0",
"sphinxcontrib-mermaid==0.9.2",
"sphinx-toolbox==4.0.0",
"docutils==0.21.2",
"myst-parser==3.0.1",
"sphinxcontrib-mermaid==1.0.0",
"nvtripy[doc_test_common]",
# Needed for guides:
"nvidia-modelopt==0.11.1",
"transformers==4.44.2",
"datasets==2.21.0",
"nvidia-modelopt==0.29.0",
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you make sure L1 passes with this change?

"transformers==4.46.2",
"datasets==3.6.0",
]
test = [
"pytest==7.1.3",
"pytest-virtualenv==1.8.0",
"pytest-profiling==1.7.0",
"pytest-cov==4.1.0",
"pytest-xdist==3.6.1",
"pytest==7.4.4",
"pytest-virtualenv==1.8.1",
"pytest-profiling==1.8.1",
"pytest-cov==6.2.1",
"pytest-xdist==3.8.0",
"pytest-benchmark==4.0.0",
"pytest-lazy-fixture==0.6.3",
"pytest-mock==3.14.0",
"pytest-mock==3.14.1",
"path.py==12.5.0",
# Triton is required for torch.compile
"triton==3.0.0",
"snakeviz==2.2.0",
"coverage==7.4.1",
"vulture==2.11",
"snakeviz==2.2.2",
"coverage==7.9.2",
"vulture==2.14",
"nvtripy[doc_test_common]",
"pytest-notebook==0.10.0",
"notebook==7.2.2",
"polygraphy==0.49.20",
"notebook==7.4.4",
"polygraphy==0.49.24",
]

[tool.black]
Expand Down
Loading