Skip to content

Commit

Permalink
Tensor[Shape] now uses dim names
Browse files Browse the repository at this point in the history
  • Loading branch information
holl- committed Feb 20, 2024
1 parent 831c705 commit 6740139
Showing 1 changed file with 3 additions and 1 deletion.
4 changes: 3 additions & 1 deletion phiml/math/magic.py
Original file line number Diff line number Diff line change
Expand Up @@ -749,6 +749,8 @@ def slicing_dict(obj, item) -> dict:
Returns:
`dict` mapping dimension names to slices.
"""
if isinstance(item, Shape):
item = item.name_list
if isinstance(item, dict):
def valid_key(key):
if callable(key):
Expand All @@ -760,7 +762,7 @@ def valid_key(key):
else:
raise ValueError(f"Illegal slicing dim: {key}. Only str and single-dim Shape and filters are allowed. While trying to slice {obj} by {item}")
return {valid_key(k): v for k, v in item.items()}
if isinstance(item, tuple):
elif isinstance(item, tuple):
if item[0] == Ellipsis:
assert len(item) - 1 == shape(obj).channel_rank
return {name: selection for name, selection in zip(channel(obj).names, item[1:])}
Expand Down

0 comments on commit 6740139

Please sign in to comment.