Merge pull request #2 from hbq1:add_pytypes

PiperOrigin-RevId: 327416525 Change-Id: I8db48986ffbacfaf5d79ed835ba9285141c9bd67
google-deepmind · Aug 19, 2020 · 83a540f · 83a540f
2 parents 8f6fe75 + 815984e
commit 83a540f
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/README.md b/README.md
@@ -20,7 +20,7 @@ or from PyPI:
 
 ## Modules Overview
 
-### Dataclass ([dataclass.py](https://github.com/deepmind/chex/blob/master/chex/_src/dataclass.p))
+### Dataclass ([dataclass.py](https://github.com/deepmind/chex/blob/master/chex/_src/dataclass.py))
 
 Dataclasses are a popular construct introduced by Python 3.7 to allows to
 easily specify typed data structures with minimal boilerplate code. They are

diff --git a/chex/_src/asserts.py b/chex/_src/asserts.py
@@ -281,7 +281,7 @@ def assert_type(
     raise AssertionError("Error in type compatibility check: " + msg + ".")
 
 
-def assert_axis_dimension(tensor, axis, expected):
+def assert_axis_dimension(tensor: Array, axis: int, expected: int):
   """Assert dimension of a specific axis of a tensor.
 
   Args: