pytc.tc.trim_panel

pytc.tc.trim_panel(tc, panel_layout, actual_a, actual_b)[source]

Strip JIT padding from the two variable axes of a tile.

When panel_size is used to stabilise XLA shapes, the raw kernel output may be larger than the true data on the padded axes. This function returns a sliced view with those axes trimmed back to their actual extents.

Parameters:

  • tc (array-like) – Raw tile of shape (Np_pad, Nq_pad, Nr_pad, Ns_pad) — only the axes listed in panel_layout may exceed the actual data.

  • panel_layout (str) –

    Which pair of axes was JIT-padded:

    • "pr" — axes 0 (p) and 2 (r)

    • "qr" — axes 1 (q) and 2 (r)

    • "ps" — axes 0 (p) and 3 (s)

  • actual_a (int) – True extent of the first padded axis (p for "pr"/"ps", q for "qr").

  • actual_b (int) – True extent of the second padded axis (r for "pr"/"qr", s for "ps").

Returns:

View of tc with the padded axes sliced to actual_a and actual_b respectively; all other axes are untouched.

Return type:

ndarray