Implement `pad` #748

jessegrabowski · 2024-05-04T16:04:37Z

Description

Implement pt.pad, following the np.pad API with feature parity.

Very preliminary draft, uploading it in this state so I can ask @ricardoV94 to look at the _linear_ramp_pad function and tell me if I'm missing something obvious related to shapes. It should follow numpy.lib.arraypad._get_linear_ramps. Also the reflection pad uses a scan, curious if we can avoid that somehow or if we think it will be no big deal (probably the 2nd).

Also I'm not sure where to put this. I put it in tensor/basic but it might be better in tensor/extra_ops?

Related Issue

Closes ENH: Implement pt.pad #743
Related to Restore fast convolution Ops, rewrites, and docs #548

Checklist

Checked that the pre-commit linting/style checks pass
Included tests that prove the fix is effective or that the new feature works
Added necessary documentation (docstrings and/or example notebooks)
If you are a pro: each commit corresponds to a relevant logical change

Type of change

ricardoV94 · 2024-05-04T17:39:10Z

What about a padding.py file?

pytensor/tensor/basic.py

jessegrabowski · 2024-05-05T01:07:02Z

Sure, I'll make a new file. It's just not my default. I agree it doesn't belong in basic.

jessegrabowski · 2024-05-11T15:22:53Z

Not quite 1:1 on numpy features but close. The more exotic padding schemes I would need more time to understand.

Still needs jax/numba overloads, but these should be very trivial.

codecov · 2024-05-11T15:43:49Z

Codecov Report

Attention: Patch coverage is 94.30380% with 18 lines in your changes missing coverage. Please review.

Project coverage is 81.48%. Comparing base (c6d85d1) to head (bbeb300).

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #748      +/-   ##
==========================================
+ Coverage   81.38%   81.48%   +0.09%     
==========================================
  Files         172      174       +2     
  Lines       46868    47166     +298     
  Branches    11423    11471      +48     
==========================================
+ Hits        38145    38434     +289     
- Misses       6542     6548       +6     
- Partials     2181     2184       +3

Files	Coverage Δ
pytensor/link/jax/dispatch/__init__.py	`100.00% <100.00%> (ø)`
pytensor/link/jax/dispatch/pad.py	`100.00% <100.00%> (ø)`
pytensor/tensor/subtensor.py	`89.31% <100.00%> (+0.11%)`	⬆️
pytensor/tensor/pad.py	`97.14% <97.14%> (ø)`
pytensor/tensor/extra_ops.py	`87.64% <80.95%> (-0.98%)`	⬇️

... and 4 files with indirect coverage changes

ricardoV94

Looks great so far, left some small suggestions.

pytensor/tensor/extra_ops.py

pytensor/tensor/pad.py

jessegrabowski · 2024-05-18T15:49:42Z

Draft of the JAX overload. Need your input on the Pad OpFromGraph. I needed some way to hang on to the keyword arguments.

It seems like there might be a difference between how JAX and numpy handle mode=mean padding, because tests pass against numpy but not against JAX. I'll investigate more carefully, but it might be a JAX bug (I doubt this padding mode is used ever)

I also think my loopy pads (symmetric, wrap) need to be redone, because they are failing a new test that arbitrarily pads every dimension of an nd input differently. So all that probably needs a re-design from the ground up.

ricardoV94 · 2024-05-20T07:43:33Z

I also think my loopy pads (symmetric, wrap) need to be redone, because they are failing a new test that arbitrarily pads every dimension of an nd input differently. So all that probably needs a re-design from the ground up.

You may need an operation per dimension

ricardoV94 · 2024-05-20T07:45:19Z

Regarding JAX do you need to implement a specific dispatch? For instance for the einsum I don't think we'll need because the OFG expression will be as good as what they do internally (since we copied it from them)

jessegrabowski · 2024-05-20T07:47:39Z

No idea on the JAX dispatch. I just assumed I should.

…tch numpy outputs

Fill out `_broadcast_inputs` docstring

This reverts commit 4cd9702.

….tensor`

pytensor/tensor/subtensor.py

pytensor/link/jax/dispatch/pad.py

jessegrabowski · 2024-07-13T15:53:25Z

Don't understand why the doctest for pad is failing

ricardoV94 · 2024-07-13T16:04:49Z

Don't understand why the doctest for pad is failing

It says there is an output that was not expected. If you have a print somewhere, you need to always test it afterward

ricardoV94 · 2024-07-13T16:06:10Z

You can also run doctest locally btw

ricardoV94 · 2024-07-13T16:07:45Z

Something like pytest --doctest-modules pytensor/tensor/pad.py --verbose

pytensor/tensor/pad.py

ricardoV94 · 2024-07-13T18:36:07Z

We should open a follow up issue for performance. With the reshape and concatenation, we're doing a lot of copies. We should see how much better it would be to have scans with set_subtensors like you tried halfway.

jessegrabowski · 2024-07-14T02:37:44Z

I kind of just want to skip the segfault test and come back to it later. I am trying to debug, but not really sure what's going on. It runs fine when NUMBA_DISABLE_NJIT flag is set. My suspicion is an out-of-range index, but I can't reproduce the error in an NB. I'd be fine just removing the numba tests all-together and considering it unsupported, but every other mode passes. So idk. I just want to get this over the finish line.

ricardoV94

This is huge!

* Add `pt.pad` * Refactor linspace, logspace, and geomspace to match numpy implementation * Add `pt.flip` * Move `flip` to `tensor/subtensor.py`, add docstring * Move `slice_at_axis` to `tensor/subtensor` and expose it in `pytensor.tensor`

ricardoV94 reviewed May 4, 2024

View reviewed changes

pytensor/tensor/basic.py Outdated Show resolved Hide resolved

ricardoV94 added enhancement New feature or request NumPy compatibility labels May 7, 2024

jessegrabowski force-pushed the pad branch from d193afa to 3186740 Compare May 11, 2024 15:21

jessegrabowski marked this pull request as ready for review May 11, 2024 15:22

jessegrabowski changed the title ~~Implement pt.pad~~ Implement pad May 11, 2024

ricardoV94 reviewed May 13, 2024

View reviewed changes

jessegrabowski added 13 commits July 7, 2024 18:35

Refactor linspace, logspace, and geomspace to match numpy implementation

91311d5

Add pt.pad

a98b8ae

Use subclassed OpFromGraph to represent pad Op

02566b6

Add test for flip

32b14d2

Address reviewer feedback

3e65827

Remove inplace argument to set_subtensor

c2b8465

Delay setting dtype of xspace Ops until after all computation to ma…

8f213a3

…tch numpy outputs

Use shape_padright instead of .reshape tricks

be6ed82

Add test for dtype kwarg on xspace Ops

9c76a8f

Save keyword arguments in Pad OpFromGraph

eeb9fa3

Add test for arbitrary padding at higher dimensions

c28faaa

First draft JAX overload

ab99a1e

Expect symbolic num argument

e93fa56

Fill out `_broadcast_inputs` docstring

jessegrabowski force-pushed the pad branch from 5499ae5 to e93fa56 Compare July 7, 2024 14:25

Split wrap_pad into separate function; eliminate use of scan

2d55a06

jessegrabowski added 4 commits July 13, 2024 21:01

Revert "Remove OpFromGraph and associated JAX dispatch"

dbda326

This reverts commit 4cd9702.

Add issue link to reflect_type error message

32ae3eb

Move flip to tensor/subtensor.py, add docstring

d543ed6

Move slice_at_axis to tensor/subtensor and expose it in `pytensor…

ba6c613

….tensor`

ricardoV94 reviewed Jul 13, 2024

View reviewed changes

pytensor/tensor/subtensor.py Outdated Show resolved Hide resolved

jessegrabowski added 4 commits July 13, 2024 22:48

Appease mypy

aa95403

Appease mypy, add docstring to pad

2a11ffa

Appease mypy, add docstring to pad

89f8cdf

Fix doctests

0f5e2ba

ricardoV94 reviewed Jul 13, 2024

View reviewed changes

pytensor/link/jax/dispatch/pad.py Outdated Show resolved Hide resolved

Fix doctests

74e39b2

Fix doctests

d169928

ricardoV94 reviewed Jul 13, 2024

View reviewed changes

pytensor/tensor/pad.py Show resolved Hide resolved

jessegrabowski added 3 commits July 14, 2024 01:12

Propagate all optional arguments to JAX

e429bab

Propagate all optional arguments to JAX

8722bfe

Appease mypy

3f8a377

jessegrabowski added 2 commits July 14, 2024 08:19

Test NUMBA backend

163b441

I love mypy

2c9f727

jessegrabowski and others added 2 commits July 14, 2024 10:46

Skip failing numba test

19ed2c0

Merge branch 'main' into pad

bbeb300

jessegrabowski requested a review from ricardoV94 July 19, 2024 12:12

ricardoV94 approved these changes Jul 19, 2024

View reviewed changes

jessegrabowski merged commit 981688c into pymc-devs:main Jul 19, 2024
58 of 59 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Implement `pad` #748

Implement `pad` #748

jessegrabowski commented May 4, 2024 •

edited

Loading

ricardoV94 commented May 4, 2024 •

edited

Loading

jessegrabowski commented May 5, 2024

jessegrabowski commented May 11, 2024

codecov bot commented May 11, 2024 •

edited

Loading

ricardoV94 left a comment

jessegrabowski commented May 18, 2024

ricardoV94 commented May 20, 2024

ricardoV94 commented May 20, 2024

jessegrabowski commented May 20, 2024

jessegrabowski commented Jul 13, 2024

ricardoV94 commented Jul 13, 2024

ricardoV94 commented Jul 13, 2024

ricardoV94 commented Jul 13, 2024 •

edited

Loading

ricardoV94 commented Jul 13, 2024

jessegrabowski commented Jul 14, 2024

ricardoV94 left a comment

Implement pad #748

Implement pad #748

Conversation

jessegrabowski commented May 4, 2024 • edited Loading

Description

Related Issue

Checklist

Type of change

ricardoV94 commented May 4, 2024 • edited Loading

jessegrabowski commented May 5, 2024

jessegrabowski commented May 11, 2024

codecov bot commented May 11, 2024 • edited Loading

Codecov Report

ricardoV94 left a comment

Choose a reason for hiding this comment

jessegrabowski commented May 18, 2024

ricardoV94 commented May 20, 2024

ricardoV94 commented May 20, 2024

jessegrabowski commented May 20, 2024

jessegrabowski commented Jul 13, 2024

ricardoV94 commented Jul 13, 2024

ricardoV94 commented Jul 13, 2024

ricardoV94 commented Jul 13, 2024 • edited Loading

ricardoV94 commented Jul 13, 2024

jessegrabowski commented Jul 14, 2024

ricardoV94 left a comment

Choose a reason for hiding this comment

Implement `pad` #748

Implement `pad` #748

jessegrabowski commented May 4, 2024 •

edited

Loading

ricardoV94 commented May 4, 2024 •

edited

Loading

codecov bot commented May 11, 2024 •

edited

Loading

ricardoV94 commented Jul 13, 2024 •

edited

Loading