João Sousa Pinto

Ownership of bounds, gradient, Hessian, and constraint matrix

3

When writing a class that encapsulates the use of this library in my codebase, I noticed that I had to store all matrices/vectors as private variables, otherwise some of the...

Add A/B comparison support

9

Hi! It would be nice to have an A/B mode for speedscope. Let me clarify what I mean by this. Suppose I have a program, and gather its initial profiling...

Unreasonable dependency on really old dm-control

1

Hi. See the error below: `The conflict is caused by: The user requested dm-control>=1.0.19 dm-acme[envs,jax,tf] 0.4.0 depends on dm-control==0.0.364896371` You'll easily reproduce this with a two-line requirements.txt file. I have...

Hi! I was able to use your Colab, but ran into some issues while trying to set things up locally (on a fresh ubuntu:latest Docker image, after having installed a...

Could you please add wheels for Python 3.11 and 3.12 in https://pypi.org/project/dm-launchpad/?

Lower tan to StableHLO instead of CHLO.

8

Fixes #23259

`jax.numpy.tan` call adds `chlo.tan` call to exported StableHLO; `stablehlo.tan` should be preferred

4

### Description Is this intended behavior? I would have expected `stablehlo` calls to be preferred to external ones when possible. Note that `stablehlo.tan` [does exist](https://openxla.org/stablehlo/spec#tan). ``` >>> import jax.numpy as...

Sparse matrix multiplication and output nse values

2

Suppose I have two sparse matrices: ``` G of type BCOO(float64[200, 80], nse=129) H of type BCOO(float64[80, 200], nse=129) ``` If I do `G @ H`, I get ``` DynamicJaxprTracer[BCOO(float64[80,...

CPU profiling (not tracing)

4

I want to get some flamegraphs from some JIT'd JAX CPU code to understand where time is being spent (in terms of my user-defined functions). My understanding (based on the...

Break up Bazel example into a separate repo to actually show GitHub Actions working

7

I think the Bazel example should be its own repository, not just a branch. The reason for this is that the GitHub Actions integration (https://github.com/wjakob/nanobind_example/blob/bazel/.github/workflows/pip.yml) on this branch doesn't actually...