Skip to content

Commit

Permalink
🐛 fix(arrayish): fix length 0 arrays (#121)
Browse files Browse the repository at this point in the history
Signed-off-by: Nathaniel Starkman <nstarman@users.noreply.github.com>
  • Loading branch information
nstarman authored Jan 21, 2025
1 parent f70fa86 commit eb24868
Show file tree
Hide file tree
Showing 2 changed files with 23 additions and 5 deletions.
26 changes: 22 additions & 4 deletions src/quaxed/experimental/_arrayish/container.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,10 +45,14 @@ class LaxLenMixin:
>>> len(x)
3
>>> x = MyArray(jnp.array(1))
>>> len(x)
0
""" # noqa: E501

def __len__(self: HasShape) -> int:
return self.shape[0]
return self.shape[0] if self.shape else 0


class NumpyLenMixin:
Expand All @@ -70,10 +74,15 @@ class NumpyLenMixin:
>>> len(x)
3
>>> x = MyArray(jnp.array(1))
>>> len(x)
0
""" # noqa: E501

def __len__(self) -> int:
return qnp.shape(self)[0]
shape = qnp.shape(self)
return shape[0] if shape else 0


# -----------------------------------------------
Expand All @@ -99,10 +108,14 @@ class LaxLengthHintMixin:
>>> x.__length_hint__()
3
>>> x = MyArray(jnp.array(0))
>>> x.__length_hint__()
0
""" # noqa: E501

def __length_hint__(self: HasShape) -> int:
return self.shape[0]
return self.shape[0] if self.shape else 0


class NumpyLengthHintMixin:
Expand All @@ -124,7 +137,12 @@ class NumpyLengthHintMixin:
>>> x.__length_hint__()
3
>>> x = MyArray(jnp.array(1))
>>> x.__length_hint__()
0
""" # noqa: E501

def __length_hint__(self) -> int:
return qnp.shape(self)[0]
shape = qnp.shape(self)
return shape[0] if shape else 0
2 changes: 1 addition & 1 deletion uv.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

0 comments on commit eb24868

Please sign in to comment.