Added code examples for ix_ function with better doc string.

[selamw1]

No description provided.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[selamw1]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

It's already signed.

[jakevdp]

Looks good, but please note that the implements decorator will overwrite any docstring you define in the function. If you want to be able to add a customized docstring, we'll have to modify that decorator (or remove it, but I believe we have tests to check that all public APIs have this decorator)

Also, it looks like this content was copied from the text within numpy repository. We cannot have verbatim copies of other open source projects as part of the JAX source code, because it violates the OSS license requirements. If you want to write a docstring for this function, I'd suggest doing it from scratch, without referencing the original NumPy docstring. Alternatively, we can put the text in jax/third_party and reference it here, although that's a bit awkward in the case of docstrings.

This licensing issue is one of the reasons we have not done this kind of docstring specialization in the past, instead relying on runtime rewrites via arguments to implements.

[selamw1]

Hi @jakevdp

Is the following modification a doable solution for inserting embedded code examples (customized to JAX) for NumPy and SciPy?

We would use lax_description for inserting examples and leave the implements decorator as it is.

@util.implements(np.ix_, lax_description="""Example:
    >>> import jax.numpy as jnp
    >>> a = jnp.arange(10).reshape(2, 5)
    >>> a
    Array([[0, 1, 2, 3, 4],
           [5, 6, 7, 8, 9]], dtype=int32)
    >>> ixgrid = jnp.ix_(jnp.array([0, 1]), jnp.array([2, 4]))
    >>> ixgrid
    (Array([[0],
           [1]], dtype=int32), Array([[2, 4]], dtype=int32))
    >>> ixgrid[0].shape, ixgrid[1].shape
    ((2, 1), (1, 2))
    >>> a[ixgrid]
    Array([[2, 4],
           [7, 9]], dtype=int32)""")
def ix_(*args: ArrayLike) -> tuple[Array, ...]:
  util.check_arraylike("ix", *args)

If not, I will do it from scratch as you suggested earlier, without referencing the original NumPy docstring.

Since the CLA has failed here, I will close and raise another PR.

[jakevdp]

Typically the example comes below the description of the function parameters. I think in this case it would be best to rewrite the docstring from scratch. I'm doing some similar work in #20941 – you can do something similar for this case.

[jakevdp]

Oh, and please don't open a new PR if possible – it's useful to keep all the conversation in one place. The failing CLA is probably due to whatever email you have attached to the commits on your branch: you can undo previous commits and redo new commits with the correct email in order to update this branch

[selamw1]

It make sene, docstring is rewritten from scratch and tried to squash the commits.

[jakevdp]

Thanks, looks good! A few formatting pieces below (you might try building the docs locally to make sure the formatting is acceptable, rather than waiting for the github CI).

Also, you might add a See also section that links to meshgrid, mgrid, and ogrid

[selamw1]

Thanks @jakevdp

'See also' section added and indentations corrected. Locally, the formatting appears as follows:

[selamw1]

Thanks @jakevdp

• Description is adjusted:
  • ndarray is modified to "Tuple of JAX arrays".
  • integer sequences is modified to integer arrays.
• New example is added.

[selamw1]

unnecessary texts and example outputs are modified/removed.

[jakevdp]

One small change requested below. Also

Please run the linter on your change (see https://jax.readthedocs.io/en/latest/contributing.html#linting-and-type-checking)

I think you'll have to add ix_ to this list to satisfy tests:

jax/tests/lax_numpy_test.py

Line 6021 in 17d0eb3

known_exceptions = {

Finally, once you've made all changes, please squash everything into a single commit; see https://jax.readthedocs.io/en/latest/contributing.html#single-change-commits-and-pull-requests for details.

[jakevdp]

Looks like there are some lint issues with whitespace. You can fix these locally by running pre-commit: https://jax.readthedocs.io/en/latest/contributing.html#linting-and-type-checking

[selamw1]

Looks like there are some lint issues with whitespace. You can fix these locally by running pre-commit: https://jax.readthedocs.io/en/latest/contributing.html#linting-and-type-checking

It seems the lint failure was below:
trim trailing whitespace.................................................Failed

removing whitespace before the following lines:
def ix_(*args: ArrayLike) -> tuple[Array, ...]:
and
util.check_arraylike("ix", *args)

resolve the issue?

running pre-commit locally prints the following error:

      Preparing metadata (pyproject.toml): finished with status 'error'
stderr:
      error: subprocess-exited-with-error
      
      × Preparing metadata (pyproject.toml) did not run successfully.
      │ exit code: 1
      ╰─> [39 lines of output]
          INFO:hatch_jupyter_builder.utils:Running jupyter-builder

[jakevdp]

Yeah, removing all trailing whitespace should fix the error. Regarding the pre-commit installation error, it's probably some kind of version incompatibility. Hard to say without seeing more of the error or logs.

[jakevdp]

It seems like some other commits snuck into your PR! It looks like you both rebased and merged, but targeted different snapshots of the main branch. I'd try fixing it like this:

$ git remote -v  # just to make sure it's clear what upstream and origin are here
origin	git@github.com:selamw1/jax.git
upstream	git@github.com:google/jax.git
$ git checkout main
$ git pull upstream main  # this should fast-forward cleanly if you haven't added commits to your local main branch
$ git checkout add_examples
$ git rebase main  # this should rebase cleanly if you haven't done anything strange
$ git log  # make sure you only have a single commit

In most cases this should work, though admittedly I don't really know what you've done to your local branch, so if you've done something really strange (like merged against branches with conflicts) then you may have to do more to fix it.

[jakevdp]

It seems like some other commits snuck into your PR! It looks like you both rebased and merged, but targeted different snapshots of the main branch in each operation. I'd try fixing it like this:

$ git remote -v  # just to make sure it's clear what upstream and origin are here
origin	git@github.com:selamw1/jax.git
upstream	git@github.com:google/jax.git
$ git checkout main
$ git pull upstream main  # this should fast-forward cleanly if you haven't added commits to your local main branch
$ git checkout add_examples
$ git rebase main  # this should rebase cleanly if you haven't done anything strange
$ git log  # make sure you only have a single commit
$ git push origin +add_examples

In most cases this should work, though admittedly I don't really know what you've done to your local branch, so if you've done something really strange (like merged against branches with conflicts) then you may have to do more to fix it.