braincore.math.ein_shape

Contents

braincore.math.ein_shape#

braincore.math.ein_shape(x, pattern)#

Parse a tensor shape to dictionary mapping axes names to their lengths.

```python # Use underscore to skip the dimension in parsing. >>> x = np.zeros([2, 3, 5, 7]) >>> ein_shape(x, ‘batch _ h w’) {‘batch’: 2, ‘h’: 5, ‘w’: 7}

# parse_shape output can be used to specify axes_lengths for other operations: >>> y = np.zeros([700]) >>> ein_rearrange(y, ‘(b c h w) -> b c h w’, **ein_shape(x, ‘b _ h w’)).shape (2, 10, 5, 7)

```

For symbolic frameworks may return symbols, not integers.

Parameters:
  • x – tensor of any supported framework

  • pattern (str) – str, space separated names for axes, underscore means skip axis

Return type:

dict

Returns:

dict, maps axes names to their lengths