shift_nth_row_n_steps package

shift_nth_row_n_steps.create_slice(ndim: int, axis_and_key: ~collections.abc.Sequence[tuple[int, int | slice | ~types.EllipsisType | None]], *, default: int | slice | ~types.EllipsisType | None | ~collections.abc.Callable[[], int | slice | ~types.EllipsisType | None] = <function <lambda>>) → tuple[int | slice | EllipsisType | None, ...]

Create a slice tuple with default values.

Parameters:

• ndim (int) – The number of dimensions.

• axis_and_key (Sequence[tuple[int, int | slice | EllipsisType | None]]) – The axis and key pair.

• default (int | slice | EllipsisType | None, optional) – The default value, by default slice(None,)

Returns:

The slice tuple.

Return type:

tuple[int | slice | EllipsisType | None, …]

shift_nth_row_n_steps.narrow(a: Array, start: int, length: int, *, axis: int) → Array

torch.narrow() in xp.

Parameters:

• a (Array) – The source array.

• start (int) – The index of the element to start from.

• length (int) – The length of the slice.

• axis (int) – The axis to narrow.

Returns:

The narrowed array.

Return type:

Array

shift_nth_row_n_steps.select(a: Array, index: int, *, axis: int) → Array

torch.select() (!= numpy.select()) in xp.

Parameters:

• a (Array) – The source array.

• index (int) – The index of the element to select.

• axis (int) – The axis to select from.

Returns:

The selected array.

Return type:

Array

shift_nth_row_n_steps.shift_nth_row_n_steps(a: Array, *, axis_row: int = -2, axis_shift: int = -1, cut_padding: bool = False, mode: Literal['fill', 'roll', 'abs'] = 'fill', fill_values: float = 0) → Array

Shifts the nth row n steps to the right.

Parameters:

• a (Array) – The source array.

• axis_row (int, optional) – The axis of the row to shift, by default -2

• axis_shift (int, optional) – The axis of the shift, by default -1

• cut_padding (bool, optional) – Whether to cut additional columns, by default False

• mode (Literal["fill", "roll", "abs"], optional) –

  The padding mode, by default “constant” - fill(padding_mode=constant) -> shift + fill

  (result[i,j] = a[i,j+n_shift*i] if j >= i else fill_values)

  • roll(padding_mode=wrap) -> shift + roll

    (a[i,j] = b[i] then result[i,j] = b[(j+n_shift*i)%len(b)])

  • abs(padding_mode=reflect) -> shift + symmetric

    (a[i,j] = b[i] then result[i,j] = b[abs(j+n_shift*i)] not implemented, do result + result.T - result * xp.eye(result.shape[-1]) instead (current behavior aims to support cut_padding = False)

• fill_values (float, optional) – The constant value to fill, by default 0 Only used when padding_mode = “constant”

Returns:

The shifted array. If the input is (…, row, …, shift, …), the output will be (…, row, …, shift + row - 1, …). […,i,…,j,…] -> […,i,…,j+i,…]

Return type:

Array

shift_nth_row_n_steps.take_slice(a: Array, start: int, end: int, *, axis: int) → Array

numpy.take() alternative using slices. (faster) similar to torch.narrow().

Parameters:

• a (Array) – The source array.

• start (int) – The index of the element to start from.

• end (int) – The index of the element to end at.

• axis (int) – The axis to take the slice from.

Returns:

The sliced array.

Return type:

Array

Submodules

shift_nth_row_n_steps.cli module

shift_nth_row_n_steps.cli.benchmark(dtype: str = <typer.models.OptionInfo object>, n_end: int = <typer.models.OptionInfo object>, n_iter: int = <typer.models.OptionInfo object>) → None

Benchmark the two implementations of the function.

Parameters:

• backend (str, optional) – The backend to use, by default typer.Option(“numpy”)

• device (str, optional) – The device to use, by default typer.Option(“cpu”)

• dtype (str, optional) – The dtype to use, by default typer.Option(“float32”)

• n_end (int, optional) – The maximum power of 2 to use, by default typer.Option(10)

• n_iter (int, optional) – The number of iterations to run, by default typer.Option(10)