Skip to content

Commit 8ee2507

Browse files
committed
Add newaxis to torch aliases
Closes data-apisgh-64
1 parent e27958b commit 8ee2507

File tree

1 file changed

+3
-2
lines changed

1 file changed

+3
-2
lines changed

array_api_compat/torch/_aliases.py

+3-2
Original file line numberDiff line numberDiff line change
@@ -144,6 +144,7 @@ def can_cast(from_: Union[Dtype, array], to: Dtype, /) -> bool:
144144

145145
# Basic renames
146146
bitwise_invert = torch.bitwise_not
147+
newaxis = None
147148

148149
# Two-arg elementwise functions
149150
# These require a wrapper to do the correct type promotion on 0-D tensors
@@ -690,8 +691,8 @@ def take(x: array, indices: array, /, *, axis: Optional[int] = None, **kwargs) -
690691
axis = 0
691692
return torch.index_select(x, axis, indices, **kwargs)
692693

693-
__all__ = ['result_type', 'can_cast', 'permute_dims', 'bitwise_invert', 'add',
694-
'atan2', 'bitwise_and', 'bitwise_left_shift', 'bitwise_or',
694+
__all__ = ['result_type', 'can_cast', 'permute_dims', 'bitwise_invert', 'newaxis',
695+
'add', 'atan2', 'bitwise_and', 'bitwise_left_shift', 'bitwise_or',
695696
'bitwise_right_shift', 'bitwise_xor', 'divide', 'equal',
696697
'floor_divide', 'greater', 'greater_equal', 'less', 'less_equal',
697698
'logaddexp', 'multiply', 'not_equal', 'pow', 'remainder',

0 commit comments

Comments
 (0)