feat: Implemented Jax front-end unwrap function #26345
Conversation
ivy-leaves
added
the
JAX Frontend
| ) | ||
| @to_ivy_arrays_and_back | ||
| def unwrap(p, discont=None, axis=-1, period=2 * ivy.pi): | ||
| p = ivy.asarray(p) |
There was a problem hiding this comment.
Is there a reason for the explicit conversion happening here?
There was a problem hiding this comment.
@bipinKrishnan - I implemented this as per the JAX source code here
| def unwrap(p, discont=None, axis=-1, period=2 * ivy.pi): | ||
| p = ivy.asarray(p) | ||
| dtype_str = str(p.dtype) | ||
| if "int" in dtype_str: |
There was a problem hiding this comment.
Instead of multiple lines of code and if-elif statements, we could use a dict mapping similar to jax here.
There was a problem hiding this comment.
@bipinKrishnan Done, replaced with dict mapping
| @@ -148,6 +148,30 @@ def _get_dtype_input_and_vectors(draw): | |||
| return dtype, vec1, vec2 | |||
|
|
|||
|
|
|||
| # unwrap | |||
| @st.composite | |||
| def _get_dtype_input_vector_axis(draw): | |||
There was a problem hiding this comment.
Helper functions should be in the respective helper modules in tests.
There was a problem hiding this comment.
@bipinKrishnan - While I was adding the function to helper module, I found an existing function that I can use. So I have used that existing function in the test instead so new helper function isn't required.
| period, | ||
| ): | ||
| dtype, x, axis = dtype_x_axis | ||
| assume(not np.any(np.isclose(x[0], 0))) |
There was a problem hiding this comment.
This should be handled while generating the values, so you could remove this line.
There was a problem hiding this comment.
@bipinKrishnan - Done this has been removed now.
|
Thank you for this PR, here is the CI results: This pull request does not result in any additional test failures. Congratulations! |
|
There's a key error issue when testing locally, can you rebase on top of main so that the CI picks up the newly added tests? |
…ted dtypes in paddle backend functions
…ted dtypes in paddle backend functions
# Conflicts: # ivy/functional/backends/paddle/searching.py # ivy/functional/backends/paddle/statistical.py
rohitkg83
force-pushed
the
jax-numpy-unwrap
branch
3 times, most recently
from
c5c354d
to
980ceba
Compare
|
@bipinKrishnan - I did some blunders during rebasing this from main. So I have created a new branch on top of latest main branch code. The PR is now available here #27050. So I am closing this PR. |
Implemented Jax numpy unwrap function and associated tests.
#25767
Close #25767
Checklist
Socials