Added code examples for ix_ function with better doc string.

[selamw1]

[google-cla[bot]]

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.

[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.

[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]

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	[email protected]:selamw1/jax.git
upstream	[email protected]: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	[email protected]:selamw1/jax.git
upstream	[email protected]: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.